<commit_before>package itpkg

import (
	"bytes"
	"fmt"
	"github.com/codegangsta/cli"
	"github.com/op/go-logging"
	"net/http"
	"os"
)

type Application struct {
	cfg *Config
}

func (p *Application) Init(env string) error {

	cfg := Config{beans: make(map[string]interface{}, 0)}
	if err := LoadLocales("locales"); err != nil {
		return err
	}

	if err := loadConfig(&cfg, fmt.Sprintf("config/%s.yml", env)); err != nil {
		return err
	}

	cfg.env = env
	p.cfg = &cfg
	if p.cfg.IsProduction() {
		bkd, err := logging.NewSyslogBackend("itpkg")
		if err != nil {
			return err
		}
		lvl := logging.AddModuleLevel(bkd)
		lvl.SetLevel(logging.INFO, "")
		logging.SetBackend(lvl)
	} else {
		logging.SetLevel(logging.DEBUG, "")
		logging.SetFormatter(logging.MustStringFormatter("%{color}%{time:15:04:05.000} %{shortfunc} ▶ %{level:.4s} %{id:03x}%{color:reset} %{message}"))
	}
	return nil
}

func (p *Application) loop(f func(en Engine)) error {
	err := p.cfg.OpenDb()
	if err != nil {
		return err
	}
	p.cfg.OpenRedis()

	if err = p.cfg.OpenCache(); err != nil {
		return err
	}

	if err = p.cfg.OpenRouter(); err != nil {
		return err
	}

	p.cfg.beans["db"] = p.cfg.db
	p.cfg.beans["cache"] = p.cfg.cache

	for _, en := range engines {
		n, v, _ := en.Info()
		Logger.Info("Load engine %s(%s)", n, v)
		en.Setup(p.cfg)
		f(en)
	}
	return nil
}

func (p *Application) Migrate() error {
	return p.loop(func(en Engine) {
		en.Migrate()
	})
}

func (p *Application) Server() error {
	if err := p.loop(func(en Engine) {
		en.Map()
		en.Mount()
	}); err != nil {
		return err
	}
	return http.ListenAndServe(fmt.Sprintf(":%d", p.cfg.Http.Port), p.cfg.router)
}

func (p *Application) Db() error {
	cmd, args := p.cfg.DbShell()
	return Shell(cmd, args...)
}

func (p *Application) Redis() error {
	cmd, args := p.cfg.RedisShell()
	return Shell(cmd, args...)
}

func (p *Application) Openssl() {
	var buf bytes.Buffer
	fmt.Fprintf(&buf,
		`
openssl genrsa -out root/root-key.pem 2048
openssl req -new -key root/root-key.pem -out root/root-req.csr -text
openssl x509 -req -in root/root-req.csr -out root/root-cert.pem -sha512 -signkey root/root-key.pem -days 3650 -text -extfile /etc/ssl/openssl.cnf -extensions v3_ca

openssl genrsa -out server/server-key.pem 2048
openssl req -new -key server/server-key.pem -out server/server-req.csr -text
openssl x509 -req -in server/server-req.csr -CA root/root-cert.pem -CAkey root/root-key.pem -CAcreateserial -days 3650 -out server/server-cert.pem -text

openssl verify -CAfile root/root-cert.pem server/server-cert.pem
openssl rsa -noout -text -in server-key.pem
openssl req -noout -text -in server-req.csr
openssl x509 -noout -text -in server-cert.pem
`)
	fmt.Fprintf(&buf, "\n")
	buf.WriteTo(os.Stdout)
}

func (p *Application) Nginx() {
	hn, _ := os.Hostname()
	wd, _ := os.Getwd()

	var buf bytes.Buffer
	fmt.Fprintf(
		&buf,
		`
upstream ksana.conf {
server http://localhost:%d fail_timeout=0;
}
`,
		p.cfg.Http.Port)
	fmt.Fprintf(
		&buf,
		`
server {
listen 443;
ssl  on;
ssl_certificate  ssl/ksana.crt;
ssl_certificate_key  ssl/ksana.key;
ssl_session_timeout  5m;
ssl_protocols  SSLv2 SSLv3 TLSv1;
ssl_ciphers  RC4:HIGH:!aNULL:!MD5;
ssl_prefer_server_ciphers  on;

client_max_body_size 4G;
keepalive_timeout 10;

server_name %s;

root %s/public;
try_files $uri $uri/index.html @ksana.conf;

location @ksana.conf {
proxy_set_header X-Forwarded-Proto https;
proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for;
proxy_set_header Host $http_host;
proxy_set_header  X-Real-IP $remote_addr;
proxy_redirect off;
proxy_pass http://ksana.conf;
# limit_req zone=one;
access_log log/ksana.access.log;
error_log log/ksana.error.log;
}

#location ^~ /assets/ {
location ~* \.(?:css|js|html|jpg|jpeg|gif|png|ico)$ {
gzip_static on;
expires max;
add_header Cache-Control public;
}

location = /50x.html {
root html;
}

location = /404.html {
root html;
}

location @503 {
error_page 405 = /system/maintenance.html;
if (-f $document_root/system/maintenance.html) {
rewrite ^(.*)$ /system/maintenance.html break;
}
rewrite ^(.*)$ /503.html break;
}

if ($request_method !~ ^(GET|HEAD|PUT|PATCH|POST|DELETE|OPTIONS)$ ){
return 405;
}

if (-f $document_root/system/maintenance.html) {
return 503;
}

location ~ \.(php|jsp|asp)$ {
return 405;
}

}
`, hn, wd)
	fmt.Fprintf(&buf, "\n")
	buf.WriteTo(os.Stdout)
}

func Run() error {

	load := func(c *cli.Context) Application {
		a := Application{}
		if err := a.Init(c.String("environment")); err != nil {
			Logger.Fatalf("error on load config:%v", err)
		}
		return a
	}

	envF := cli.StringFlag{
		Name:   "environment, e",
		Value:  "development",
		Usage:  "can be production, development, etc...",
		EnvVar: "ITPKG_ENV",
	}

	app := cli.NewApp()
	app.Name = "itpkg"
	app.Usage = "IT-PACKAGE"
	app.Version = "v20150529"
	app.Commands = []cli.Command{
		{
			Name:    "server",
			Aliases: []string{"s"},
			Usage:   "Start the ITPKG server",
			Flags: []cli.Flag{
				envF,
				cli.IntFlag{
					Name:   "port, p",
					Value:  3000,
					Usage:  "listen port",
					EnvVar: "PORT",
				},
			},
			Action: func(c *cli.Context) {
				a := load(c)
				if e := a.Server(); e != nil {
					Logger.Fatalf("Error on start server: %v", e)
				}
			},
		},
		{
			Name:    "dbconsole",
			Aliases: []string{"db"},
			Usage:   "Start a console for the database",
			Flags:   []cli.Flag{envF},
			Action: func(c *cli.Context) {
				a := load(c)
				a.Db()
			},
		},
		{
			Name:    "redis",
			Aliases: []string{"r"},
			Usage:   "Start a console for the redis",
			Flags:   []cli.Flag{envF},
			Action: func(c *cli.Context) {
				a := load(c)
				a.Redis()
			},
		},
		{
			Name:    "nginx",
			Aliases: []string{"n"},
			Usage:   "Nginx config file demo",
			Flags:   []cli.Flag{envF},
			Action: func(c *cli.Context) {
				a := load(c)
				a.Nginx()
			},
		},
		{
			Name:    "openssl",
			Aliases: []string{"ssl"},
			Usage:   "Openssl certs command demo",
			Flags:   []cli.Flag{envF},
			Action: func(c *cli.Context) {
				a := Application{}
				a.Openssl()
			},
		},
		{
			Name:    "migrate",
			Aliases: []string{"m"},
			Usage:   "Migrate the database",
			Flags:   []cli.Flag{envF},
			Action: func(c *cli.Context) {
				a := load(c)
				a.Migrate()
			},
		},
	}

	return app.Run(os.Args)
}
<commit_msg>show routes<commit_after>package itpkg

import (
	"bytes"
	"fmt"
	"github.com/codegangsta/cli"
	"github.com/gin-gonic/gin"
	"github.com/op/go-logging"
	"net/http"
	"os"
)

type Application struct {
	cfg *Config
}

func (p *Application) Init(env string) error {

	cfg := Config{beans: make(map[string]interface{}, 0)}
	if err := LoadLocales("locales"); err != nil {
		return err
	}

	if err := loadConfig(&cfg, fmt.Sprintf("config/%s.yml", env)); err != nil {
		return err
	}

	cfg.env = env
	p.cfg = &cfg
	if p.cfg.IsProduction() {
		bkd, err := logging.NewSyslogBackend("itpkg")
		if err != nil {
			return err
		}
		lvl := logging.AddModuleLevel(bkd)
		lvl.SetLevel(logging.INFO, "")
		logging.SetBackend(lvl)
	} else {
		logging.SetLevel(logging.DEBUG, "")
		logging.SetFormatter(logging.MustStringFormatter("%{color}%{time:15:04:05.000} %{shortfunc} ▶ %{level:.4s} %{id:03x}%{color:reset} %{message}"))
	}
	return nil
}

func (p *Application) loop(f func(en Engine)) error {
	err := p.cfg.OpenDb()
	if err != nil {
		return err
	}
	p.cfg.OpenRedis()

	if err = p.cfg.OpenCache(); err != nil {
		return err
	}

	if err = p.cfg.OpenRouter(); err != nil {
		return err
	}

	p.cfg.beans["db"] = p.cfg.db
	p.cfg.beans["cache"] = p.cfg.cache

	for _, en := range engines {
		n, v, _ := en.Info()
		Logger.Info("Load engine %s(%s)", n, v)
		en.Setup(p.cfg)
		f(en)
	}
	return nil
}

func (p *Application) Migrate() error {
	return p.loop(func(en Engine) {
		en.Migrate()
	})
}

func (p *Application) Server() error {
	if err := p.loop(func(en Engine) {
		en.Map()
		en.Mount()
	}); err != nil {
		return err
	}
	return http.ListenAndServe(fmt.Sprintf(":%d", p.cfg.Http.Port), p.cfg.router)
}

func (p *Application) Routes() error {
	if err := p.loop(func(en Engine) {
		gin.SetMode(gin.DebugMode)
		en.Mount()
	}); err != nil {
		return err
	}

	return nil
}

func (p *Application) Db() error {
	cmd, args := p.cfg.DbShell()
	return Shell(cmd, args...)
}

func (p *Application) Redis() error {
	cmd, args := p.cfg.RedisShell()
	return Shell(cmd, args...)
}

func (p *Application) Openssl() {
	var buf bytes.Buffer
	fmt.Fprintf(&buf,
		`
openssl genrsa -out root/root-key.pem 2048
openssl req -new -key root/root-key.pem -out root/root-req.csr -text
openssl x509 -req -in root/root-req.csr -out root/root-cert.pem -sha512 -signkey root/root-key.pem -days 3650 -text -extfile /etc/ssl/openssl.cnf -extensions v3_ca

openssl genrsa -out server/server-key.pem 2048
openssl req -new -key server/server-key.pem -out server/server-req.csr -text
openssl x509 -req -in server/server-req.csr -CA root/root-cert.pem -CAkey root/root-key.pem -CAcreateserial -days 3650 -out server/server-cert.pem -text

openssl verify -CAfile root/root-cert.pem server/server-cert.pem
openssl rsa -noout -text -in server-key.pem
openssl req -noout -text -in server-req.csr
openssl x509 -noout -text -in server-cert.pem
`)
	fmt.Fprintf(&buf, "\n")
	buf.WriteTo(os.Stdout)
}

func (p *Application) Nginx() {
	hn, _ := os.Hostname()
	wd, _ := os.Getwd()

	var buf bytes.Buffer
	fmt.Fprintf(
		&buf,
		`
upstream ksana.conf {
server http://localhost:%d fail_timeout=0;
}
`,
		p.cfg.Http.Port)
	fmt.Fprintf(
		&buf,
		`
server {
listen 443;
ssl  on;
ssl_certificate  ssl/ksana.crt;
ssl_certificate_key  ssl/ksana.key;
ssl_session_timeout  5m;
ssl_protocols  SSLv2 SSLv3 TLSv1;
ssl_ciphers  RC4:HIGH:!aNULL:!MD5;
ssl_prefer_server_ciphers  on;

client_max_body_size 4G;
keepalive_timeout 10;

server_name %s;

root %s/public;
try_files $uri $uri/index.html @ksana.conf;

location @ksana.conf {
proxy_set_header X-Forwarded-Proto https;
proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for;
proxy_set_header Host $http_host;
proxy_set_header  X-Real-IP $remote_addr;
proxy_redirect off;
proxy_pass http://ksana.conf;
# limit_req zone=one;
access_log log/ksana.access.log;
error_log log/ksana.error.log;
}

#location ^~ /assets/ {
location ~* \.(?:css|js|html|jpg|jpeg|gif|png|ico)$ {
gzip_static on;
expires max;
add_header Cache-Control public;
}

location = /50x.html {
root html;
}

location = /404.html {
root html;
}

location @503 {
error_page 405 = /system/maintenance.html;
if (-f $document_root/system/maintenance.html) {
rewrite ^(.*)$ /system/maintenance.html break;
}
rewrite ^(.*)$ /503.html break;
}

if ($request_method !~ ^(GET|HEAD|PUT|PATCH|POST|DELETE|OPTIONS)$ ){
return 405;
}

if (-f $document_root/system/maintenance.html) {
return 503;
}

location ~ \.(php|jsp|asp)$ {
return 405;
}

}
`, hn, wd)
	fmt.Fprintf(&buf, "\n")
	buf.WriteTo(os.Stdout)
}

func Run() error {

	load := func(c *cli.Context) Application {
		a := Application{}
		if err := a.Init(c.String("environment")); err != nil {
			Logger.Fatalf("error on load config:%v", err)
		}
		return a
	}

	envF := cli.StringFlag{
		Name:   "environment, e",
		Value:  "development",
		Usage:  "can be production, development, etc...",
		EnvVar: "ITPKG_ENV",
	}

	app := cli.NewApp()
	app.Name = "itpkg"
	app.Usage = "IT-PACKAGE"
	app.Version = "v20150529"
	app.Commands = []cli.Command{
		{
			Name:    "server",
			Aliases: []string{"s"},
			Usage:   "Start the ITPKG server",
			Flags: []cli.Flag{
				envF,
				cli.IntFlag{
					Name:   "port, p",
					Value:  3000,
					Usage:  "listen port",
					EnvVar: "PORT",
				},
			},
			Action: func(c *cli.Context) {
				a := load(c)
				if e := a.Server(); e != nil {
					Logger.Fatalf("Error on start server: %v", e)
				}
			},
		},
		{
			Name:    "dbconsole",
			Aliases: []string{"db"},
			Usage:   "Start a console for the database",
			Flags:   []cli.Flag{envF},
			Action: func(c *cli.Context) {
				a := load(c)
				a.Db()
			},
		},
		{
			Name:    "routes",
			Aliases: []string{"ro"},
			Usage:   "Print out all defined routes in match order, with names",
			Flags:   []cli.Flag{envF},
			Action: func(c *cli.Context) {
				a := load(c)
				a.Routes()
			},
		},
		{
			Name:    "redis",
			Aliases: []string{"re"},
			Usage:   "Start a console for the redis",
			Flags:   []cli.Flag{envF},
			Action: func(c *cli.Context) {
				a := load(c)
				a.Redis()
			},
		},
		{
			Name:    "nginx",
			Aliases: []string{"n"},
			Usage:   "Nginx config file demo",
			Flags:   []cli.Flag{envF},
			Action: func(c *cli.Context) {
				a := load(c)
				a.Nginx()
			},
		},
		{
			Name:    "openssl",
			Aliases: []string{"ssl"},
			Usage:   "Openssl certs command demo",
			Flags:   []cli.Flag{envF},
			Action: func(c *cli.Context) {
				a := Application{}
				a.Openssl()
			},
		},
		{
			Name:    "migrate",
			Aliases: []string{"m"},
			Usage:   "Migrate the database",
			Flags:   []cli.Flag{envF},
			Action: func(c *cli.Context) {
				a := load(c)
				a.Migrate()
			},
		},
	}

	return app.Run(os.Args)
}
<|endoftext|>


<commit_before>// +build integration

package integration

import (
	log "github.com/Sirupsen/logrus"
	"github.com/intelsdi-x/swan/pkg/executor"
	"github.com/intelsdi-x/swan/pkg/workloads/memcached"
	. "github.com/smartystreets/goconvey/convey"
	"io/ioutil"
	"os"
	"strings"
	"testing"
	"time"
)

const (
	netstatCommand       = "echo stats | nc -w 1 127.0.0.1 11211"
	defaultMemcachedPath = "workloads/data_caching/memcached/memcached-1.4.25/build/memcached"
	swanPkg              = "github.com/intelsdi-x/swan"
)

// TestMemcachedWithExecutor is an integration test with local executor.
// See README for setup items.
func TestMemcachedWithExecutor(t *testing.T) {
	log.SetLevel(log.DebugLevel)

	// Get optional custom Memcached path from MEMCACHED_PATH.
	memcachedPath := os.Getenv("MEMCACHED_BIN")

	if memcachedPath == "" {
		// If custom path does not exists use default path for built memcached.
		pathSlice := []string{os.Getenv("GOPATH"), "src", swanPkg, defaultMemcachedPath}
		memcachedPath = strings.Join(pathSlice, "/")
	}

	Convey("While using Local Shell in Memcached launcher", t, func() {
		l := executor.NewLocal()
		memcachedLauncher := memcached.New(
			l, memcached.DefaultMemcachedConfig(memcachedPath))

		Convey("When memcached is launched", func() {
			// NOTE: It is needed for memcached to have default port available.
			task, err := memcachedLauncher.Launch()

			Convey("There should be no error", func() {
				stopErr := task.Stop()

				So(err, ShouldBeNil)
				So(stopErr, ShouldBeNil)
			})

			Convey("Wait 1 second for memcached to init", func() {
				isTerminated := task.Wait(1 * time.Second)

				Convey("Memcached should be still running", func() {
					stopErr := task.Stop()

					// NOTE: Here you will be failing if the memcached
					// can start because it needs to have default port available.
					So(isTerminated, ShouldBeFalse)

					So(stopErr, ShouldBeNil)
				})

				Convey("When we check the memcached endpoint for stats after 1 second", func() {

					netstatTask, netstatErr := l.Execute(netstatCommand)

					Convey("There should be no error", func() {
						task.Stop()
						netstatTask.Stop()

						So(netstatErr, ShouldBeNil)

					})

					Convey("When we wait for netstat ", func() {
						netstatTask.Wait(0)

						Convey("The netstat task should be terminated, the task status should be 0"+
							" and output resultes with a STAT information", func() {
							netstatTaskState, netstatTaskStatus := netstatTask.Status()

							task.Stop()
							netstatTask.Stop()

							So(netstatTaskState, ShouldEqual, executor.TERMINATED)
							So(netstatTaskStatus.ExitCode, ShouldEqual, 0)

							stdoutReader, stdoutErr := netstatTask.Stdout()
							So(stdoutErr, ShouldBeNil)
							So(stdoutReader, ShouldNotBeNil)

							data, readErr := ioutil.ReadAll(stdoutReader)
							So(readErr, ShouldBeNil)
							So(string(data[:]), ShouldStartWith, "STAT")
						})
					})

					if netstatTask != nil {
						netstatTask.Clean()
					}
				})

				Convey("When we stop the memcached task", func() {
					err := task.Stop()

					Convey("There should be no error", func() {
						So(err, ShouldBeNil)
					})

					Convey("The task should be terminated and the task status "+
						"should be -1 or 0", func() {
						taskState, taskStatus := task.Status()
						So(taskState, ShouldEqual, executor.TERMINATED)
						// Memcached on CentOS returns 0 (successful code) after SIGTERM.
						So(taskStatus.ExitCode, ShouldBeIn, -1, 0)
					})
				})
			})

			if task != nil {
				task.Clean()
			}
		})
	})
}
<commit_msg>Issues addressed.<commit_after>// +build integration

package integration

import (
	log "github.com/Sirupsen/logrus"
	"github.com/intelsdi-x/swan/pkg/executor"
	"github.com/intelsdi-x/swan/pkg/workloads/memcached"
	. "github.com/smartystreets/goconvey/convey"
	"io/ioutil"
	"os"
	"strings"
	"testing"
	"time"
)

const (
	netstatCommand       = "echo stats | nc -w 1 127.0.0.1 11211"
	defaultMemcachedPath = "workloads/data_caching/memcached/memcached-1.4.25/build/memcached"
	swanPkg              = "github.com/intelsdi-x/swan"
)

// TestMemcachedWithExecutor is an integration test with local executor.
// See README for setup items.
func TestMemcachedWithExecutor(t *testing.T) {
	log.SetLevel(log.DebugLevel)

	// Get optional custom Memcached path from MEMCACHED_PATH.
	memcachedPath := os.Getenv("MEMCACHED_BIN")

	if memcachedPath == "" {
		// If custom path does not exists use default path for built memcached.
		pathSlice := []string{os.Getenv("GOPATH"), "src", swanPkg, defaultMemcachedPath}
		memcachedPath = strings.Join(pathSlice, "/")
	}

	Convey("While using Local Shell in Memcached launcher", t, func() {
		l := executor.NewLocal()
		memcachedLauncher := memcached.New(
			l, memcached.DefaultMemcachedConfig(memcachedPath))

		Convey("When memcached is launched", func() {
			// NOTE: It is needed for memcached to have default port available.
			task, err := memcachedLauncher.Launch()
			if task != nil {
				defer task.Stop()
				defer task.Clean()
				defer task.EraseOutput()
			}

			Convey("There should be no error", func() {
				stopErr := task.Stop()

				So(err, ShouldBeNil)
				So(stopErr, ShouldBeNil)
			})

			Convey("Wait 1 second for memcached to init", func() {
				isTerminated := task.Wait(1 * time.Second)

				Convey("Memcached should be still running", func() {
					stopErr := task.Stop()

					// NOTE: Here you will be failing if the memcached
					// can start because it needs to have default port available.
					So(isTerminated, ShouldBeFalse)

					So(stopErr, ShouldBeNil)
				})

				Convey("When we check the memcached endpoint for stats after 1 second", func() {

					netstatTask, netstatErr := l.Execute(netstatCommand)
					if netstatTask != nil {
						defer netstatTask.Stop()
						defer netstatTask.Clean()
						defer netstatTask.EraseOutput()
					}
					Convey("There should be no error", func() {
						task.Stop()
						netstatTask.Stop()

						So(netstatErr, ShouldBeNil)

					})

					Convey("When we wait for netstat ", func() {
						netstatTask.Wait(0)

						Convey("The netstat task should be terminated, the task status should be 0"+
							" and output resultes with a STAT information", func() {
							netstatTaskState, netstatTaskStatus := netstatTask.Status()

							So(netstatTaskState, ShouldEqual, executor.TERMINATED)
							So(netstatTaskStatus.ExitCode, ShouldEqual, 0)

							stdoutReader, stdoutErr := netstatTask.Stdout()
							So(stdoutErr, ShouldBeNil)
							So(stdoutReader, ShouldNotBeNil)

							data, readErr := ioutil.ReadAll(stdoutReader)
							So(readErr, ShouldBeNil)
							So(string(data[:]), ShouldStartWith, "STAT")
						})
					})
				})

				Convey("When we stop the memcached task", func() {
					err := task.Stop()

					Convey("There should be no error", func() {
						So(err, ShouldBeNil)
					})

					Convey("The task should be terminated and the task status "+
						"should be -1 or 0", func() {
						taskState, taskStatus := task.Status()
						So(taskState, ShouldEqual, executor.TERMINATED)
						// Memcached on CentOS returns 0 (successful code) after SIGTERM.
						So(taskStatus.ExitCode, ShouldBeIn, -1, 0)
					})
				})
			})
		})
	})
}
<|endoftext|>


<commit_before>// Copyright 2016 Red Hat, Inc, and individual contributors.
//
// 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 main

import (
	"sort"
	"strings"

	"github.com/spf13/cobra"
	"github.com/spf13/pflag"
)

type CommandCompleter struct {
	Root   *cobra.Command
	Finder ResourceFinder
}

func (cc *CommandCompleter) Do(line []rune, pos int) (newLine [][]rune, offset int) {
	cmd, args := cc.Root, []string{}
	word := string(line[:pos])
	lastSpace := strings.LastIndex(string(line[:pos]), " ") + 1
	lastComma := strings.LastIndex(string(line[:pos]), ",") + 1
	if lastSpace > 0 {
		if lastComma > lastSpace {
			word = word[lastComma:pos]
		} else {
			word = word[lastSpace:pos]
		}
		var err error
		cmd, args, err = cc.Root.Find(strings.Split(string(line), " "))
		if err != nil {
			return
		}
	}
	for _, completion := range cc.completions(word, cmd, args) {
		if len(word) >= len(completion) {
			if len(word) == len(completion) {
				newLine = append(newLine, []rune{' '})
			} else {
				newLine = append(newLine, []rune(completion))
			}
			offset = len(completion)
		} else {
			newLine = append(newLine, []rune(completion)[len(word):])
			offset = len(word)
		}

	}
	return
}

func (cc *CommandCompleter) completions(prefix string, cmd *cobra.Command, args []string) []string {
	candidates := []string{}
	if strings.HasPrefix(prefix, "-") {
		candidates = flags(cmd)
	} else {
		candidates = subCommands(cmd)
		if len(candidates) == 0 {
			switch cmd.Name() {
			case "logs", "attach", "exec", "port-forward":
				candidates = cc.resources("pods")
			case "rolling-update":
				candidates = cc.resources("rc")
			case "cordon", "uncordon", "drain":
				candidates = cc.resources("node")
			default:
				if t := resourceType(args); len(t) > 0 {
					candidates = cc.resources(t)
				} else {
					candidates = ResourceTypes(cmd)
				}
			}
		}
	}
	return complete(prefix, candidates)
}

func (cc *CommandCompleter) resources(resourceType string) []string {

	resources, err := cc.Finder.Lookup([]string{resourceType})
	ret := []string{}

	if err == nil {
		for _, r := range resources {
			ret = append(ret, r.name)
		}
	}

	return ret
}

func complete(prefix string, candidates []string) (results []string) {
	for _, s := range candidates {
		if strings.HasPrefix(s, prefix) {
			results = append(results, s)
		}
	}
	return
}

func subCommands(cmd *cobra.Command) []string {
	prefixes := make([]string, len(cmd.Commands()))
	for i, c := range cmd.Commands() {
		prefixes[i] = c.Name() + " "
	}
	return prefixes
}

func flags(cmd *cobra.Command) []string {
	flags := []string{}
	fn := func(f *pflag.Flag) {
		if len(f.Deprecated) > 0 || f.Hidden {
			return
		}
		flag := "--" + f.Name
		if len(f.NoOptDefVal) == 0 {
			flag += "="
		} else {
			flag += " "
		}
		flags = append(flags, flag)
	}
	cmd.NonInheritedFlags().VisitAll(fn)
	cmd.InheritedFlags().VisitAll(fn)
	return flags
}

func ResourceTypes(cmd *cobra.Command) []string {
	args := cmd.ValidArgs
	sort.Strings(args)
	return args
}

// resourceType returns the resource type identified in the args,
// which could be a comma-delimited list of multiple types
func resourceType(args []string) string {
	x := []string{}
	for _, s := range args {
		if !strings.HasPrefix(s, "-") {
			x = append(x, s)
		}
	}
	if len(x) > 1 {
		return x[0]
	} else {
		return ""
	}
}
<commit_msg>Don't tab-complete deprecated cmds to fix #16<commit_after>// Copyright 2016 Red Hat, Inc, and individual contributors.
//
// 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 main

import (
	"sort"
	"strings"

	"github.com/spf13/cobra"
	"github.com/spf13/pflag"
)

type CommandCompleter struct {
	Root   *cobra.Command
	Finder ResourceFinder
}

func (cc *CommandCompleter) Do(line []rune, pos int) (newLine [][]rune, offset int) {
	cmd, args := cc.Root, []string{}
	word := string(line[:pos])
	lastSpace := strings.LastIndex(string(line[:pos]), " ") + 1
	lastComma := strings.LastIndex(string(line[:pos]), ",") + 1
	if lastSpace > 0 {
		if lastComma > lastSpace {
			word = word[lastComma:pos]
		} else {
			word = word[lastSpace:pos]
		}
		var err error
		cmd, args, err = cc.Root.Find(strings.Split(string(line), " "))
		if err != nil {
			return
		}
	}
	for _, completion := range cc.completions(word, cmd, args) {
		if len(word) >= len(completion) {
			if len(word) == len(completion) {
				newLine = append(newLine, []rune{' '})
			} else {
				newLine = append(newLine, []rune(completion))
			}
			offset = len(completion)
		} else {
			newLine = append(newLine, []rune(completion)[len(word):])
			offset = len(word)
		}

	}
	return
}

func (cc *CommandCompleter) completions(prefix string, cmd *cobra.Command, args []string) []string {
	candidates := []string{}
	if strings.HasPrefix(prefix, "-") {
		candidates = flags(cmd)
	} else {
		candidates = subCommands(cmd)
		if len(candidates) == 0 {
			switch cmd.Name() {
			case "logs", "attach", "exec", "port-forward":
				candidates = cc.resources("pods")
			case "rolling-update":
				candidates = cc.resources("rc")
			case "cordon", "uncordon", "drain":
				candidates = cc.resources("node")
			default:
				if t := resourceType(args); len(t) > 0 {
					candidates = cc.resources(t)
				} else {
					candidates = ResourceTypes(cmd)
				}
			}
		}
	}
	return complete(prefix, candidates)
}

func (cc *CommandCompleter) resources(resourceType string) []string {

	resources, err := cc.Finder.Lookup([]string{resourceType})
	ret := []string{}

	if err == nil {
		for _, r := range resources {
			ret = append(ret, r.name)
		}
	}

	return ret
}

func complete(prefix string, candidates []string) (results []string) {
	for _, s := range candidates {
		if strings.HasPrefix(s, prefix) {
			results = append(results, s)
		}
	}
	return
}

func subCommands(cmd *cobra.Command) []string {
	prefixes := make([]string, 0, len(cmd.Commands()))
	for _, c := range cmd.Commands() {
		if c.IsAvailableCommand() {
			prefixes = append(prefixes, c.Name()+" ")
		}
	}
	return prefixes
}

func flags(cmd *cobra.Command) []string {
	flags := []string{}
	fn := func(f *pflag.Flag) {
		if len(f.Deprecated) > 0 || f.Hidden {
			return
		}
		flag := "--" + f.Name
		if len(f.NoOptDefVal) == 0 {
			flag += "="
		} else {
			flag += " "
		}
		flags = append(flags, flag)
	}
	cmd.NonInheritedFlags().VisitAll(fn)
	cmd.InheritedFlags().VisitAll(fn)
	return flags
}

func ResourceTypes(cmd *cobra.Command) []string {
	args := cmd.ValidArgs
	sort.Strings(args)
	return args
}

// resourceType returns the resource type identified in the args,
// which could be a comma-delimited list of multiple types
func resourceType(args []string) string {
	x := []string{}
	for _, s := range args {
		if !strings.HasPrefix(s, "-") {
			x = append(x, s)
		}
	}
	if len(x) > 1 {
		return x[0]
	} else {
		return ""
	}
}
<|endoftext|>


<commit_before>/*
 * Ferret
 * Copyright (c) 2016 Yieldbot, Inc.
 * For the full copyright and license information, please view the LICENSE.txt file.
 */

// Package api provides REST API functionality
package api

import (
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"os"

	"github.com/yieldbot/ferret/search"
	"golang.org/x/net/context"
)

var (
	// Logger is the logger
	Logger *log.Logger

	listenPort = "3030"
)

func init() {
	Logger = log.New(os.Stderr, "", log.LstdFlags)

	if e := os.Getenv("FERRET_LISTEN_PORT"); e != "" {
		listenPort = e
	}
}

// Listen initializes HTTP handlers and listens for the requests
func Listen() {
	// Init handlers
	http.HandleFunc("/search", SearchHandler)

	// Listen
	log.Printf("listening on %s", listenPort)
	if err := http.ListenAndServe(fmt.Sprintf(":%s", listenPort), nil); err != nil {
		log.Fatal(err)
	}
}

// SearchHandler is the handler for search route
func SearchHandler(w http.ResponseWriter, req *http.Request) {

	// Search
	q := search.Query{
		Provider: req.URL.Query().Get("provider"),
		Keyword:  req.URL.Query().Get("keyword"),
		Page:     search.ParsePage(req.URL.Query().Get("page")),
		Timeout:  search.ParseTimeout(req.URL.Query().Get("timeout")),
	}
	q, err := search.Do(context.Background(), q)
	if err != nil {
		http.Error(w, err.Error(), q.HTTPStatus)
		return
	}

	// Prepare data
	var data []byte
	if len(q.Results) > 0 {
		if req.URL.Query().Get("output") == "pretty" {
			data, err = json.MarshalIndent(q.Results, "", "  ")
		} else {
			data, err = json.Marshal(q.Results)
		}
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
	}

	// Return
	cb := req.URL.Query().Get("callback")
	if cb != "" {
		w.Header().Set("Content-Type", "application/javascript")
		fmt.Fprintf(w, "%s(%s)", cb, data)
	} else {
		w.Header().Set("Content-Type", "application/json")
		w.Write(data)
	}
}
<commit_msg>Remove logger for now<commit_after>/*
 * Ferret
 * Copyright (c) 2016 Yieldbot, Inc.
 * For the full copyright and license information, please view the LICENSE.txt file.
 */

// Package api provides REST API functionality
package api

import (
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"os"

	"github.com/yieldbot/ferret/search"
	"golang.org/x/net/context"
)

var (
	listenPort = "3030"
)

func init() {
	if e := os.Getenv("FERRET_LISTEN_PORT"); e != "" {
		listenPort = e
	}
}

// Listen initializes HTTP handlers and listens for the requests
func Listen() {
	// Init handlers
	http.HandleFunc("/search", SearchHandler)

	// Listen
	log.Printf("listening on %s", listenPort)
	if err := http.ListenAndServe(fmt.Sprintf(":%s", listenPort), nil); err != nil {
		log.Fatal(err)
	}
}

// SearchHandler is the handler for search route
func SearchHandler(w http.ResponseWriter, req *http.Request) {

	// Search
	q := search.Query{
		Provider: req.URL.Query().Get("provider"),
		Keyword:  req.URL.Query().Get("keyword"),
		Page:     search.ParsePage(req.URL.Query().Get("page")),
		Timeout:  search.ParseTimeout(req.URL.Query().Get("timeout")),
	}
	q, err := search.Do(context.Background(), q)
	if err != nil {
		http.Error(w, err.Error(), q.HTTPStatus)
		return
	}

	// Prepare data
	var data []byte
	if len(q.Results) > 0 {
		if req.URL.Query().Get("output") == "pretty" {
			data, err = json.MarshalIndent(q.Results, "", "  ")
		} else {
			data, err = json.Marshal(q.Results)
		}
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
	}

	// Return
	cb := req.URL.Query().Get("callback")
	if cb != "" {
		w.Header().Set("Content-Type", "application/javascript")
		fmt.Fprintf(w, "%s(%s)", cb, data)
	} else {
		w.Header().Set("Content-Type", "application/json")
		w.Write(data)
	}
}
<|endoftext|>


<commit_before>package grpc

import (
	"context"
	"fmt"
	"net"
	"reflect"
	"sync"
	"testing"
	"time"

	"github.com/golang/protobuf/proto"
	"github.com/google/cloudprober/logger"
	"github.com/google/cloudprober/metrics"
	"github.com/google/cloudprober/probes/options"
	probepb "github.com/google/cloudprober/probes/proto"
	grpcpb "github.com/google/cloudprober/servers/grpc/proto"
	spb "github.com/google/cloudprober/servers/grpc/proto"
	"github.com/google/cloudprober/targets"
	"github.com/google/cloudprober/targets/endpoint"
	"github.com/google/cloudprober/targets/resolver"
	"google.golang.org/grpc"
)

var once sync.Once
var srvAddr string
var baseProbeConf = `
name: "grpc"
type: GRPC
targets {
	host_names: "%s"
}
interval_msec: 1000
timeout_msec: %d
grpc_probe {
	%s
	num_conns: %d
	connect_timeout_msec: 2000
}
`

func probeCfg(tgts, cred string, timeout, numConns int) (*probepb.ProbeDef, error) {
	conf := fmt.Sprintf(baseProbeConf, tgts, timeout, cred, numConns)
	cfg := &probepb.ProbeDef{}
	err := proto.UnmarshalText(conf, cfg)
	return cfg, err
}

type Server struct {
	delay time.Duration
	msg   []byte
}

// Echo reflects back the incoming message.
// TODO: return error if EchoMessage is greater than maxMsgSize.
func (s *Server) Echo(ctx context.Context, req *spb.EchoMessage) (*spb.EchoMessage, error) {
	if s.delay > 0 {
		time.Sleep(s.delay)
	}
	return req, nil
}

// BlobRead returns a blob of data.
func (s *Server) BlobRead(ctx context.Context, req *spb.BlobReadRequest) (*spb.BlobReadResponse, error) {
	return &spb.BlobReadResponse{
		Blob: s.msg[0:req.GetSize()],
	}, nil
}

// ServerStatus returns the current server status.
func (s *Server) ServerStatus(ctx context.Context, req *spb.StatusRequest) (*spb.StatusResponse, error) {
	return &spb.StatusResponse{
		UptimeUs: proto.Int64(42),
	}, nil
}

// BlobWrite returns the size of blob in the WriteRequest. It does not operate
// on the blob.
func (s *Server) BlobWrite(ctx context.Context, req *spb.BlobWriteRequest) (*spb.BlobWriteResponse, error) {
	return &spb.BlobWriteResponse{
		Size: proto.Int32(int32(len(req.Blob))),
	}, nil
}

// globalGRPCServer sets up runconfig and returns a gRPC server.
func globalGRPCServer() (string, error) {
	var err error
	once.Do(func() {
		var ln net.Listener
		ln, err = net.Listen("tcp", "localhost:0")
		if err != nil {
			return
		}
		grpcSrv := grpc.NewServer()
		srv := &Server{delay: time.Second / 2, msg: make([]byte, 1024)}
		grpcpb.RegisterProberServer(grpcSrv, srv)
		go grpcSrv.Serve(ln)
		srvAddr = ln.Addr().String()
		time.Sleep(time.Second * 2)
	})
	return srvAddr, err
}

// TestGRPCSuccess tests probe output on success.
// 2 connections, 1 probe/sec/conn, stats exported every 5 sec
// 	=> 5-10 results/interval. Test looks for minimum of 7 results.
func TestGRPCSuccess(t *testing.T) {
	addr, err := globalGRPCServer()
	if err != nil {
		t.Fatalf("Error initializing global config: %v", err)
	}
	cfg, err := probeCfg(addr, "", 1000, 2)
	if err != nil {
		t.Fatalf("Error unmarshalling config: %v", err)
	}
	l := &logger.Logger{}
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	iters := 5
	statsExportInterval := time.Duration(iters) * time.Second

	probeOpts := &options.Options{
		Targets:             targets.StaticTargets(addr),
		Timeout:             time.Second * 1,
		Interval:            time.Second * 1,
		ProbeConf:           cfg.GetGrpcProbe(),
		Logger:              l,
		StatsExportInterval: statsExportInterval,
		LogMetrics:          func(em *metrics.EventMetrics) {},
	}
	p := &Probe{}
	p.Init("grpc-success", probeOpts)
	dataChan := make(chan *metrics.EventMetrics, 5)
	go p.Start(ctx, dataChan)
	time.Sleep(statsExportInterval * 2)
	found := false
	expectedLabels := map[string]string{
		"ptype": "grpc",
		"dst":   addr,
		"probe": "grpc-success",
	}

	for i := 0; i < 2; i++ {
		select {
		case em := <-dataChan:
			t.Logf("Probe results: %v", em.String())
			total := em.Metric("total").(*metrics.Int)
			success := em.Metric("success").(*metrics.Int)
			expect := int64(iters) + 2
			if total.Int64() < expect || success.Int64() < expect {
				t.Errorf("Got total=%d success=%d, expecting at least %d for each", total.Int64(), success.Int64(), expect)
			}
			gotLabels := make(map[string]string)
			for _, k := range em.LabelsKeys() {
				gotLabels[k] = em.Label(k)
			}
			if !reflect.DeepEqual(gotLabels, expectedLabels) {
				t.Errorf("Unexpected labels: got: %v, expected: %v", gotLabels, expectedLabels)
			}
			found = true
		default:
			time.Sleep(time.Second)
		}
	}
	if !found {
		t.Errorf("No probe results found")
	}
}

// TestConnectFailures forces the use of GCE credentials resulting in connect
// failures and stats exported once every connect timeout.
// 2 connections, 0.5 connect attempt/sec/conn, stats exported every 6 sec
//  => 3 - 6 connect errors/sec. Test looks for minimum of 4 attempts.
func TestConnectFailures(t *testing.T) {
	addr, err := globalGRPCServer()
	if err != nil {
		t.Fatalf("Error initializing global config: %v", err)
	}
	cfg, err := probeCfg(addr, "alts_config: {}", 1000, 2)
	if err != nil {
		t.Fatalf("Error unmarshalling config: %v", err)
	}
	l := &logger.Logger{}
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	iters := 6
	statsExportInterval := time.Duration(iters) * time.Second

	probeOpts := &options.Options{
		Targets:             targets.StaticTargets(addr),
		Timeout:             time.Second * 1,
		Interval:            time.Second * 1,
		ProbeConf:           cfg.GetGrpcProbe(),
		Logger:              l,
		StatsExportInterval: statsExportInterval,
		LogMetrics:          func(em *metrics.EventMetrics) {},
	}
	p := &Probe{}
	p.Init("grpc-connectfail", probeOpts)
	dataChan := make(chan *metrics.EventMetrics, 5)
	go p.Start(ctx, dataChan)
	time.Sleep(statsExportInterval * 2)
	found := false
	for i := 0; i < 2; i++ {
		select {
		case em := <-dataChan:
			t.Logf("Probe results: %v", em.String())
			total := em.Metric("total").(*metrics.Int)
			success := em.Metric("success").(*metrics.Int)
			connectErrs := em.Metric("connecterrors").(*metrics.Int)
			expect := int64(iters/2) + 1
			if success.Int64() > 0 {
				t.Errorf("Got %d probe successes, want all failures", success.Int64())
			}
			if total.Int64() < expect || connectErrs.Int64() < expect {
				t.Errorf("Got total=%d connectErrs=%d, expecting at least %d for each", total.Int64(), connectErrs.Int64(), expect)
			}
			found = true
		default:
			time.Sleep(time.Second)
		}
	}
	if !found {
		t.Errorf("No probe results found")
	}
}

func TestProbeTimeouts(t *testing.T) {
	addr, err := globalGRPCServer()
	if err != nil {
		t.Fatalf("Error initializing global config: %v", err)
	}
	cfg, err := probeCfg(addr, "", 1000, 1)
	if err != nil {
		t.Fatalf("Error unmarshalling config: %v", err)
	}
	l := &logger.Logger{}
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	iters := 5
	statsExportInterval := time.Duration(iters) * time.Second

	probeOpts := &options.Options{
		Targets:             targets.StaticTargets(addr),
		Timeout:             time.Millisecond * 100,
		Interval:            time.Second * 1,
		ProbeConf:           cfg.GetGrpcProbe(),
		Logger:              l,
		LatencyUnit:         time.Millisecond,
		StatsExportInterval: statsExportInterval,
		LogMetrics:          func(em *metrics.EventMetrics) {},
	}
	p := &Probe{}
	p.Init("grpc-reqtimeout", probeOpts)
	dataChan := make(chan *metrics.EventMetrics, 5)
	go p.Start(ctx, dataChan)
	time.Sleep(statsExportInterval * 2)
	found := false
	for i := 0; i < 2; i++ {
		select {
		case em := <-dataChan:
			t.Logf("Probe results: %v", em.String())
			total := em.Metric("total").(*metrics.Int)
			success := em.Metric("success").(*metrics.Int)
			expect := int64(iters) - 1
			if success.Int64() > 0 {
				t.Errorf("Got %d probe successes, want all failures", success.Int64())
			}
			if total.Int64() < expect {
				t.Errorf("Got total=%d, want at least %d", total.Int64(), expect)
			}
			found = true
		default:
			time.Sleep(time.Second)
		}
	}
	if !found {
		t.Errorf("No probe results found")
	}
}

type testTargets struct {
	r *resolver.Resolver

	start        time.Time
	startTargets []endpoint.Endpoint

	switchDur   time.Duration
	nextTargets []endpoint.Endpoint
}

func newTargets(startTargets, nextTargets []endpoint.Endpoint, switchDur time.Duration) targets.Targets {
	return &testTargets{r: resolver.New(), startTargets: startTargets, nextTargets: nextTargets, start: time.Now(), switchDur: switchDur}
}

func (t *testTargets) ListEndpoints() []endpoint.Endpoint {
	if time.Since(t.start) > t.switchDur {
		return t.nextTargets
	}
	return t.startTargets
}

func (t *testTargets) List() []string {
	var targets []string
	for _, ep := range t.ListEndpoints() {
		targets = append(targets, ep.Name)
	}
	return targets
}

func (t *testTargets) Resolve(name string, ipVer int) (net.IP, error) {
	return t.r.Resolve(name, ipVer)
}
<commit_msg>Force probe failures by connecting to discard port<commit_after>package grpc

import (
	"context"
	"fmt"
	"net"
	"reflect"
	"sync"
	"testing"
	"time"

	"github.com/golang/protobuf/proto"
	"github.com/google/cloudprober/logger"
	"github.com/google/cloudprober/metrics"
	"github.com/google/cloudprober/probes/options"
	probepb "github.com/google/cloudprober/probes/proto"
	grpcpb "github.com/google/cloudprober/servers/grpc/proto"
	spb "github.com/google/cloudprober/servers/grpc/proto"
	"github.com/google/cloudprober/targets"
	"github.com/google/cloudprober/targets/endpoint"
	"github.com/google/cloudprober/targets/resolver"
	"google.golang.org/grpc"
)

var once sync.Once
var srvAddr string
var baseProbeConf = `
name: "grpc"
type: GRPC
targets {
	host_names: "%s"
}
interval_msec: 1000
timeout_msec: %d
grpc_probe {
	%s
	num_conns: %d
	connect_timeout_msec: 2000
}
`

func probeCfg(tgts, cred string, timeout, numConns int) (*probepb.ProbeDef, error) {
	conf := fmt.Sprintf(baseProbeConf, tgts, timeout, cred, numConns)
	cfg := &probepb.ProbeDef{}
	err := proto.UnmarshalText(conf, cfg)
	return cfg, err
}

type Server struct {
	delay time.Duration
	msg   []byte
}

// Echo reflects back the incoming message.
// TODO: return error if EchoMessage is greater than maxMsgSize.
func (s *Server) Echo(ctx context.Context, req *spb.EchoMessage) (*spb.EchoMessage, error) {
	if s.delay > 0 {
		time.Sleep(s.delay)
	}
	return req, nil
}

// BlobRead returns a blob of data.
func (s *Server) BlobRead(ctx context.Context, req *spb.BlobReadRequest) (*spb.BlobReadResponse, error) {
	return &spb.BlobReadResponse{
		Blob: s.msg[0:req.GetSize()],
	}, nil
}

// ServerStatus returns the current server status.
func (s *Server) ServerStatus(ctx context.Context, req *spb.StatusRequest) (*spb.StatusResponse, error) {
	return &spb.StatusResponse{
		UptimeUs: proto.Int64(42),
	}, nil
}

// BlobWrite returns the size of blob in the WriteRequest. It does not operate
// on the blob.
func (s *Server) BlobWrite(ctx context.Context, req *spb.BlobWriteRequest) (*spb.BlobWriteResponse, error) {
	return &spb.BlobWriteResponse{
		Size: proto.Int32(int32(len(req.Blob))),
	}, nil
}

// globalGRPCServer sets up runconfig and returns a gRPC server.
func globalGRPCServer() (string, error) {
	var err error
	once.Do(func() {
		var ln net.Listener
		ln, err = net.Listen("tcp", "localhost:0")
		if err != nil {
			return
		}
		grpcSrv := grpc.NewServer()
		srv := &Server{delay: time.Second / 2, msg: make([]byte, 1024)}
		grpcpb.RegisterProberServer(grpcSrv, srv)
		go grpcSrv.Serve(ln)
		srvAddr = ln.Addr().String()
		time.Sleep(time.Second * 2)
	})
	return srvAddr, err
}

// TestGRPCSuccess tests probe output on success.
// 2 connections, 1 probe/sec/conn, stats exported every 5 sec
// 	=> 5-10 results/interval. Test looks for minimum of 7 results.
func TestGRPCSuccess(t *testing.T) {
	addr, err := globalGRPCServer()
	if err != nil {
		t.Fatalf("Error initializing global config: %v", err)
	}
	cfg, err := probeCfg(addr, "", 1000, 2)
	if err != nil {
		t.Fatalf("Error unmarshalling config: %v", err)
	}
	l := &logger.Logger{}
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	iters := 5
	statsExportInterval := time.Duration(iters) * time.Second

	probeOpts := &options.Options{
		Targets:             targets.StaticTargets(addr),
		Timeout:             time.Second * 1,
		Interval:            time.Second * 1,
		ProbeConf:           cfg.GetGrpcProbe(),
		Logger:              l,
		StatsExportInterval: statsExportInterval,
		LogMetrics:          func(em *metrics.EventMetrics) {},
	}
	p := &Probe{}
	p.Init("grpc-success", probeOpts)
	dataChan := make(chan *metrics.EventMetrics, 5)
	go p.Start(ctx, dataChan)
	time.Sleep(statsExportInterval * 2)
	found := false
	expectedLabels := map[string]string{
		"ptype": "grpc",
		"dst":   addr,
		"probe": "grpc-success",
	}

	for i := 0; i < 2; i++ {
		select {
		case em := <-dataChan:
			t.Logf("Probe results: %v", em.String())
			total := em.Metric("total").(*metrics.Int)
			success := em.Metric("success").(*metrics.Int)
			expect := int64(iters) + 2
			if total.Int64() < expect || success.Int64() < expect {
				t.Errorf("Got total=%d success=%d, expecting at least %d for each", total.Int64(), success.Int64(), expect)
			}
			gotLabels := make(map[string]string)
			for _, k := range em.LabelsKeys() {
				gotLabels[k] = em.Label(k)
			}
			if !reflect.DeepEqual(gotLabels, expectedLabels) {
				t.Errorf("Unexpected labels: got: %v, expected: %v", gotLabels, expectedLabels)
			}
			found = true
		default:
			time.Sleep(time.Second)
		}
	}
	if !found {
		t.Errorf("No probe results found")
	}
}

// TestConnectFailures attempts to connect to localhost:9 (discard port) and
// checks that stats are exported once every connect timeout.
// 2 connections, 0.5 connect attempt/sec/conn, stats exported every 6 sec
//  => 3 - 6 connect errors/sec. Test looks for minimum of 4 attempts.
func TestConnectFailures(t *testing.T) {
	addr := "localhost:9"
	cfg, err := probeCfg(addr, "", 1000, 2)
	if err != nil {
		t.Fatalf("Error unmarshalling config: %v", err)
	}
	l := &logger.Logger{}
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	iters := 6
	statsExportInterval := time.Duration(iters) * time.Second

	probeOpts := &options.Options{
		Targets:             targets.StaticTargets(addr),
		Timeout:             time.Second * 1,
		Interval:            time.Second * 1,
		ProbeConf:           cfg.GetGrpcProbe(),
		Logger:              l,
		StatsExportInterval: statsExportInterval,
		LogMetrics:          func(em *metrics.EventMetrics) {},
	}
	p := &Probe{}
	p.Init("grpc-connectfail", probeOpts)
	dataChan := make(chan *metrics.EventMetrics, 5)
	go p.Start(ctx, dataChan)
	time.Sleep(statsExportInterval * 2)
	found := false
	for i := 0; i < 2; i++ {
		select {
		case em := <-dataChan:
			t.Logf("Probe results: %v", em.String())
			total := em.Metric("total").(*metrics.Int)
			success := em.Metric("success").(*metrics.Int)
			connectErrs := em.Metric("connecterrors").(*metrics.Int)
			expect := int64(iters/2) + 1
			if success.Int64() > 0 {
				t.Errorf("Got %d probe successes, want all failures", success.Int64())
			}
			if total.Int64() < expect || connectErrs.Int64() < expect {
				t.Errorf("Got total=%d connectErrs=%d, expecting at least %d for each", total.Int64(), connectErrs.Int64(), expect)
			}
			found = true
		default:
			time.Sleep(time.Second)
		}
	}
	if !found {
		t.Errorf("No probe results found")
	}
}

func TestProbeTimeouts(t *testing.T) {
	addr, err := globalGRPCServer()
	if err != nil {
		t.Fatalf("Error initializing global config: %v", err)
	}
	cfg, err := probeCfg(addr, "", 1000, 1)
	if err != nil {
		t.Fatalf("Error unmarshalling config: %v", err)
	}
	l := &logger.Logger{}
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	iters := 5
	statsExportInterval := time.Duration(iters) * time.Second

	probeOpts := &options.Options{
		Targets:             targets.StaticTargets(addr),
		Timeout:             time.Millisecond * 100,
		Interval:            time.Second * 1,
		ProbeConf:           cfg.GetGrpcProbe(),
		Logger:              l,
		LatencyUnit:         time.Millisecond,
		StatsExportInterval: statsExportInterval,
		LogMetrics:          func(em *metrics.EventMetrics) {},
	}
	p := &Probe{}
	p.Init("grpc-reqtimeout", probeOpts)
	dataChan := make(chan *metrics.EventMetrics, 5)
	go p.Start(ctx, dataChan)
	time.Sleep(statsExportInterval * 2)
	found := false
	for i := 0; i < 2; i++ {
		select {
		case em := <-dataChan:
			t.Logf("Probe results: %v", em.String())
			total := em.Metric("total").(*metrics.Int)
			success := em.Metric("success").(*metrics.Int)
			expect := int64(iters) - 1
			if success.Int64() > 0 {
				t.Errorf("Got %d probe successes, want all failures", success.Int64())
			}
			if total.Int64() < expect {
				t.Errorf("Got total=%d, want at least %d", total.Int64(), expect)
			}
			found = true
		default:
			time.Sleep(time.Second)
		}
	}
	if !found {
		t.Errorf("No probe results found")
	}
}

type testTargets struct {
	r *resolver.Resolver

	start        time.Time
	startTargets []endpoint.Endpoint

	switchDur   time.Duration
	nextTargets []endpoint.Endpoint
}

func newTargets(startTargets, nextTargets []endpoint.Endpoint, switchDur time.Duration) targets.Targets {
	return &testTargets{r: resolver.New(), startTargets: startTargets, nextTargets: nextTargets, start: time.Now(), switchDur: switchDur}
}

func (t *testTargets) ListEndpoints() []endpoint.Endpoint {
	if time.Since(t.start) > t.switchDur {
		return t.nextTargets
	}
	return t.startTargets
}

func (t *testTargets) List() []string {
	var targets []string
	for _, ep := range t.ListEndpoints() {
		targets = append(targets, ep.Name)
	}
	return targets
}

func (t *testTargets) Resolve(name string, ipVer int) (net.IP, error) {
	return t.r.Resolve(name, ipVer)
}
<|endoftext|>


<commit_before>// Discordgo - Discord bindings for Go
// Available at https://github.com/bwmarrin/discordgo

// Copyright 2015-2016 Bruce Marriner <bruce@sqls.net>.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// This file contains code related to the Message struct

package discordgo

import (
	"fmt"
	"strings"
)

// A Message stores all data related to a specific Discord message.
type Message struct {
	ID              string        `json:"id"`
	Author          *User         `json:"author"`
	Content         string        `json:"content"`
	Attachments     []*Attachment `json:"attachments"`
	Tts             bool          `json:"tts"`
	Embeds          []*Embed      `json:"embeds"`
	Timestamp       string        `json:"timestamp"`
	MentionEveryone bool          `json:"mention_everyone"`
	EditedTimestamp string        `json:"edited_timestamp"`
	Mentions        []*User       `json:"mentions"`
	ChannelID       string        `json:"channel_id"`
}

// An Attachment stores data for message attachments.
type Attachment struct { //TODO figure this out
}

// An Embed stores data for message embeds.
type Embed struct { // TODO figure this out
}

// ContentWithMentionsReplaced will replace all @<id> mentions with the
// username of the mention.
func (m *Message) ContentWithMentionsReplaced() string {
	content := m.Content
	for _, user := range m.Mentions {
		content = strings.Replace(content, fmt.Sprintf("<@%s>", user.ID),
			fmt.Sprintf("@%s", user.Username), -1)
	}
	return content
}
<commit_msg>Add attachments fields.<commit_after>// Discordgo - Discord bindings for Go
// Available at https://github.com/bwmarrin/discordgo

// Copyright 2015-2016 Bruce Marriner <bruce@sqls.net>.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// This file contains code related to the Message struct

package discordgo

import (
	"fmt"
	"strings"
)

// A Message stores all data related to a specific Discord message.
type Message struct {
	ID              string        `json:"id"`
	Author          *User         `json:"author"`
	Content         string        `json:"content"`
	Attachments     []*Attachment `json:"attachments"`
	Tts             bool          `json:"tts"`
	Embeds          []*Embed      `json:"embeds"`
	Timestamp       string        `json:"timestamp"`
	MentionEveryone bool          `json:"mention_everyone"`
	EditedTimestamp string        `json:"edited_timestamp"`
	Mentions        []*User       `json:"mentions"`
	ChannelID       string        `json:"channel_id"`
}

// An Attachment stores data for message attachments.
type Attachment struct {
	Width    int    `json:"width"`
	URL      string `json:"url"`
	Size     int    `json:"size"`
	ProxyURL string `json:"proxy_url"`
	ID       string `json:"id"`
	Height   int    `json:"height"`
	Filename string `json:"filename"`
}

// An Embed stores data for message embeds.
type Embed struct { // TODO figure this out
}

// ContentWithMentionsReplaced will replace all @<id> mentions with the
// username of the mention.
func (m *Message) ContentWithMentionsReplaced() string {
	content := m.Content
	for _, user := range m.Mentions {
		content = strings.Replace(content, fmt.Sprintf("<@%s>", user.ID),
			fmt.Sprintf("@%s", user.Username), -1)
	}
	return content
}
<|endoftext|>


<commit_before>package msgqueue

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

	"github.com/vmihailenco/msgpack"
)

var ErrDuplicate = errors.New("queue: message with such name already exists")

// Message is used to create and retrieve messages from a queue.
type Message struct {
	// SQS/IronMQ message id.
	Id string

	// Optional name for the message. Messages with the same name
	// are processed only once.
	Name string

	// Delay specifies the duration the queue must wait
	// before executing the message.
	Delay time.Duration

	// Function args passed to the handler.
	Args []interface{}

	// Text representation of the Args.
	Body    string
	bodyErr error

	// SQS/IronMQ reservation id that is used to release/delete the message..
	ReservationId string

	// The number of times the message has been reserved or released.
	ReservedCount int

	Err error
}

func NewMessage(args ...interface{}) *Message {
	return &Message{
		Args: args,
	}
}

func (m *Message) String() string {
	return fmt.Sprintf(
		"Message<Id=%q Name=%q ReservedCount=%d>",
		m.Id, m.Name, m.ReservedCount,
	)
}

// SetDelayName sets delay and generates message name from the args.
func (m *Message) SetDelayName(delay time.Duration, args ...interface{}) {
	m.Name = argsName(append(args, timeSlot(delay)))
	m.Delay = delay
	// Some random delay to better distribute the load.
	m.Delay += time.Duration(rand.Intn(5)+1) * time.Second
}

func (m *Message) EncodeArgs() (string, error) {
	if m.bodyErr != nil {
		return "", m.bodyErr
	}
	if m.Body != "" {
		return m.Body, nil
	}
	m.Body, m.bodyErr = encodeArgs(m.Args)
	return m.Body, m.bodyErr
}

func timeSlot(resolution time.Duration) int64 {
	resolutionInSeconds := int64(resolution / time.Second)
	if resolutionInSeconds <= 0 {
		return 0
	}
	return time.Now().Unix() / resolutionInSeconds
}

func argsName(args []interface{}) string {
	b, _ := msgpack.Marshal(args...)
	return string(b)
}
<commit_msg>Allow zero delay in SetDelayName<commit_after>package msgqueue

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

	"github.com/vmihailenco/msgpack"
)

var ErrDuplicate = errors.New("queue: message with such name already exists")

// Message is used to create and retrieve messages from a queue.
type Message struct {
	// SQS/IronMQ message id.
	Id string

	// Optional name for the message. Messages with the same name
	// are processed only once.
	Name string

	// Delay specifies the duration the queue must wait
	// before executing the message.
	Delay time.Duration

	// Function args passed to the handler.
	Args []interface{}

	// Text representation of the Args.
	Body    string
	bodyErr error

	// SQS/IronMQ reservation id that is used to release/delete the message..
	ReservationId string

	// The number of times the message has been reserved or released.
	ReservedCount int

	Err error
}

func NewMessage(args ...interface{}) *Message {
	return &Message{
		Args: args,
	}
}

func (m *Message) String() string {
	return fmt.Sprintf(
		"Message<Id=%q Name=%q ReservedCount=%d>",
		m.Id, m.Name, m.ReservedCount,
	)
}

// SetDelayName sets delay and generates message name from the args.
func (m *Message) SetDelayName(delay time.Duration, args ...interface{}) {
	m.Name = argsName(append(args, timeSlot(delay)))
	m.Delay = delay
	if m.Delay > 3*time.Second {
		// Some random delay to better distribute the load.
		m.Delay += time.Duration(rand.Intn(5)+1) * time.Second
	}
}

func (m *Message) EncodeArgs() (string, error) {
	if m.bodyErr != nil {
		return "", m.bodyErr
	}
	if m.Body != "" {
		return m.Body, nil
	}
	m.Body, m.bodyErr = encodeArgs(m.Args)
	return m.Body, m.bodyErr
}

func timeSlot(resolution time.Duration) int64 {
	resolutionInSeconds := int64(resolution / time.Second)
	if resolutionInSeconds <= 0 {
		return 0
	}
	return time.Now().Unix() / resolutionInSeconds
}

func argsName(args []interface{}) string {
	b, _ := msgpack.Marshal(args...)
	return string(b)
}
<|endoftext|>


<commit_before>package avs

import (
	"encoding/json"
	"fmt"
	"reflect"
	"strings"
	"time"
)

// An interface that represents both raw Message objects and more specifically
// typed ones. Usually, values of this interface are used with a type switch:
//	switch d := typedMessage.(type) {
//	case *Speak:
//		fmt.Printf("We got a spoken response in format %s.\n", d.Payload.Format)
//	}
//
type TypedMessage interface {
	GetMessage() *Message
	Typed() TypedMessage
}

// A general structure for contexts, events and directives.
type Message struct {
	Header  map[string]string `json:"header"`
	Payload json.RawMessage   `json:"payload,omitempty"`
}

// Creates a Message suited for being used as a context value.
func NewContext(namespace, name string) *Message {
	return &Message{
		Header: map[string]string{
			"namespace": namespace,
			"name":      name,
		},
		Payload: nil,
	}
}

// Creates a Message suited for being used as an event value.
func NewEvent(namespace, name, messageId, dialogRequestId string) *Message {
	m := &Message{
		Header: map[string]string{
			"namespace": namespace,
			"name":      name,
			"messageId": messageId,
		},
		Payload: nil,
	}
	if dialogRequestId != "" {
		m.Header["dialogRequestId"] = dialogRequestId
	}
	return m
}

// Returns a pointer to the underlying Message object.
func (m *Message) GetMessage() *Message {
	return m
}

// Returns the namespace and name as a single string.
func (m *Message) String() string {
	return fmt.Sprintf("%s.%s", m.Header["namespace"], m.Header["name"])
}

// Returns a more specific type for this message.
func (m *Message) Typed() TypedMessage {
	switch m.String() {
	case "Alerts.DeleteAlert":
		return fill(new(DeleteAlert), m)
	case "Alerts.SetAlert":
		return fill(new(SetAlert), m)
	case "AudioPlayer.ClearQueue":
		return fill(new(ClearQueue), m)
	case "AudioPlayer.Play":
		return fill(new(Play), m)
	case "AudioPlayer.Stop":
		return fill(new(Stop), m)
	case "SpeechRecognizer.ExpectSpeech":
		return fill(new(ExpectSpeech), m)
	case "SpeechSynthesizer.Speak":
		return fill(new(Speak), m)
	case "System.Exception":
		return fill(new(Exception), m)
	case "System.SynchronizeState":
		return fill(new(SynchronizeState), m)
	default:
		return m
	}
}

// The AlertEnteredBackground event.
type AlertEnteredBackground struct {
	*Message
	Payload struct {
		Token string `json:"token"`
	} `json:"payload"`
}

func NewAlertEnteredBackground(messageId, token string) *AlertEnteredBackground {
	m := new(AlertEnteredBackground)
	m.Message = NewEvent("Alerts", "AlertEnteredBackground", messageId, "")
	m.Payload.Token = token
	return m
}

// The AlertEnteredForeground event.
type AlertEnteredForeground struct {
	*Message
	Payload struct {
		Token string `json:"token"`
	} `json:"payload"`
}

func NewAlertEnteredForeground(messageId, token string) *AlertEnteredForeground {
	m := new(AlertEnteredForeground)
	m.Message = NewEvent("Alerts", "AlertEnteredForeground", messageId, "")
	m.Payload.Token = token
	return m
}

// The AlertStarted event.
type AlertStarted struct {
	*Message
	Payload struct {
		Token string `json:"token"`
	} `json:"payload"`
}

func NewAlertStarted(messageId, token string) *AlertStarted {
	m := new(AlertStarted)
	m.Message = NewEvent("Alerts", "AlertStarted", messageId, "")
	m.Payload.Token = token
	return m
}

// The AlertStopped event.
type AlertStopped struct {
	*Message
	Payload struct {
		Token string `json:"token"`
	} `json:"payload"`
}

func NewAlertStopped(messageId, token string) *AlertStopped {
	m := new(AlertStopped)
	m.Message = NewEvent("Alerts", "AlertStopped", messageId, "")
	m.Payload.Token = token
	return m
}

// The ClearQueue directive.
type ClearQueue struct {
	*Message
	Payload struct {
		ClearBehavior ClearBehavior `json:"clearBehavior"`
	} `json:"payload"`
}

// The DeleteAlert directive.
type DeleteAlert struct {
	*Message
	Payload struct {
		Token string `json:"token"`
	} `json:"payload"`
}

// The DeleteAlertFailed event.
type DeleteAlertFailed struct {
	*Message
	Payload struct {
		Token string `json:"token"`
	} `json:"payload"`
}

func NewDeleteAlertFailed(messageId, token string) *DeleteAlertFailed {
	m := new(DeleteAlertFailed)
	m.Message = NewEvent("Alerts", "DeleteAlertFailed", messageId, "")
	m.Payload.Token = token
	return m
}

// The DeleteAlertSucceeded event.
type DeleteAlertSucceeded struct {
	*Message
	Payload struct {
		Token string `json:"token"`
	} `json:"payload"`
}

func NewDeleteAlertSucceeded(messageId, token string) *DeleteAlertSucceeded {
	m := new(DeleteAlertSucceeded)
	m.Message = NewEvent("Alerts", "DeleteAlertSucceeded", messageId, "")
	m.Payload.Token = token
	return m
}

// The Exception message.
type Exception struct {
	*Message
	Payload struct {
		Code        string `json:"code"`
		Description string `json:"description"`
	} `json:"payload"`
}

func (m *Exception) Error() string {
	return fmt.Sprintf("%s: %s", m.Payload.Code, m.Payload.Description)
}

// The ExceptionEncountered event.
type ExceptionEncountered struct {
	*Message
	Payload struct {
		UnparsedDirective string `json:"unparsedDirective"`
		Error             struct {
			Type    string `json:"type"`
			Message string `json:"message"`
		} `json:"error"`
	} `json:"payload"`
}

func NewExceptionEncountered(messageId, directive, errorType, errorMessage string) *ExceptionEncountered {
	m := new(ExceptionEncountered)
	m.Message = NewEvent("System", "ExceptionEncountered", messageId, "")
	m.Payload.UnparsedDirective = directive
	m.Payload.Error.Type = errorType
	m.Payload.Error.Message = errorMessage
	return m
}

// The ExpectSpeech directive.
type ExpectSpeech struct {
	*Message
	Payload struct {
		TimeoutInMilliseconds float64 `json:"timeoutInMilliseconds"`
	} `json:"payload"`
}

func (m *ExpectSpeech) Timeout() time.Duration {
	return time.Duration(m.Payload.TimeoutInMilliseconds) * time.Millisecond
}

// The ExpectSpeechTimedOut event.
type ExpectSpeechTimedOut struct {
	*Message
	Payload struct{} `json:"payload"`
}

func NewExpectSpeechTimedOut(messageId string) *ExpectSpeechTimedOut {
	m := new(ExpectSpeechTimedOut)
	m.Message = NewEvent("SpeechRecognizer", "ExpectSpeechTimedOut", messageId, "")
	return m
}

// The Play directive.
type Play struct {
	*Message
	Payload struct {
		AudioItem    AudioItem    `json:"audioItem"`
		PlayBehavior PlayBehavior `json:"playBehavior"`
	} `json:"payload"`
}

func (m *Play) DialogRequestId() string {
	return m.Header["dialogRequestId"]
}

func (m *Play) MessageId() string {
	return m.Header["messageId"]
}

// The PlaybackState context.
type PlaybackState struct {
	*Message
	Payload struct {
		Token                string         `json:"token"`
		OffsetInMilliseconds float64        `json:"offsetInMilliseconds"`
		PlayerActivity       PlayerActivity `json:"playerActivity"`
	} `json:"payload"`
}

func NewPlaybackState(token string, offset time.Duration, activity PlayerActivity) *PlaybackState {
	m := new(PlaybackState)
	m.Message = NewContext("AudioPlayer", "PlaybackState")
	m.Payload.OffsetInMilliseconds = offset.Seconds() * 1000
	m.Payload.PlayerActivity = activity
	m.Payload.Token = token
	return m
}

func (m *PlaybackState) Offset() time.Duration {
	return time.Duration(m.Payload.OffsetInMilliseconds) * time.Millisecond
}

// The Recognize event.
type Recognize struct {
	*Message
	Payload struct {
		Profile string `json:"profile"`
		Format  string `json:"format"`
	} `json:"payload"`
}

func NewRecognize(messageId, dialogRequestId string) *Recognize {
	m := new(Recognize)
	m.Message = NewEvent("SpeechRecognizer", "Recognize", messageId, dialogRequestId)
	m.Payload.Format = "AUDIO_L16_RATE_16000_CHANNELS_1"
	m.Payload.Profile = "CLOSE_TALK"
	return m
}

// The SetAlert directive.
type SetAlert struct {
	*Message
	Payload struct {
		Token         string    `json:"token"`
		Type          AlertType `json:"type"`
		ScheduledTime time.Time `json:"scheduledTime"`
	} `json:"payload"`
}

// The SetAlertFailed event.
type SetAlertFailed struct {
	*Message
	Payload struct {
		Token string `json:"token"`
	} `json:"payload"`
}

func NewSetAlertFailed(messageId, token string) *SetAlertFailed {
	m := new(SetAlertFailed)
	m.Message = NewEvent("Alerts", "SetAlertFailed", messageId, "")
	m.Payload.Token = token
	return m
}

// The SetAlertSucceeded event.
type SetAlertSucceeded struct {
	*Message
	Payload struct {
		Token string `json:"token"`
	} `json:"payload"`
}

func NewSetAlertSucceeded(messageId, token string) *SetAlertSucceeded {
	m := new(SetAlertSucceeded)
	m.Message = NewEvent("Alerts", "SetAlertSucceeded", messageId, "")
	m.Payload.Token = token
	return m
}

// The Speak directive.
type Speak struct {
	*Message
	Payload struct {
		Format string
		URL    string
	} `json:"payload"`
}

func (m *Speak) ContentId() string {
	if !strings.HasPrefix(m.Payload.URL, "cid:") {
		return ""
	}
	return m.Payload.URL[4:]
}

// The ResetUserInactivity directive.
type ResetUserInactivity struct {
	*Message
	Payload struct{} `json:"payload"`
}

// The Stop directive.
type Stop struct {
	*Message
	Payload struct{} `json:"payload"`
}

// The SynchronizeState event.
type SynchronizeState struct {
	*Message
	Payload struct{} `json:"payload"`
}

func NewSynchronizeState(messageId string) *SynchronizeState {
	m := new(SynchronizeState)
	m.Message = NewEvent("System", "SynchronizeState", messageId, "")
	return m
}

// The UserInactivityReport event.
type UserInactivityReport struct {
	*Message
	Payload struct {
		InactiveTimeInSeconds float64 `json:"inactiveTimeInSeconds"`
	} `json:"payload"`
}

func NewUserInactivityReport(messageId string, inactiveTime time.Duration) *UserInactivityReport {
	m := new(UserInactivityReport)
	m.Message = NewEvent("System", "UserInactivityReport", messageId, "")
	m.Payload.InactiveTimeInSeconds = inactiveTime.Seconds()
	return m
}

// Convenience method to set up an empty typed message object from a raw Message.
func fill(dst TypedMessage, src *Message) TypedMessage {
	v := reflect.ValueOf(dst).Elem()
	v.FieldByName("Message").Set(reflect.ValueOf(src))
	payload := v.FieldByName("Payload")
	if payload.Kind() != reflect.Struct {
		return dst
	}
	json.Unmarshal(src.Payload, payload.Addr().Interface())
	return dst
}
<commit_msg>Don't use time.Time until Amazon fixes timestamp format<commit_after>package avs

import (
	"encoding/json"
	"fmt"
	"reflect"
	"strings"
	"time"
)

// An interface that represents both raw Message objects and more specifically
// typed ones. Usually, values of this interface are used with a type switch:
//	switch d := typedMessage.(type) {
//	case *Speak:
//		fmt.Printf("We got a spoken response in format %s.\n", d.Payload.Format)
//	}
//
type TypedMessage interface {
	GetMessage() *Message
	Typed() TypedMessage
}

// A general structure for contexts, events and directives.
type Message struct {
	Header  map[string]string `json:"header"`
	Payload json.RawMessage   `json:"payload,omitempty"`
}

// Creates a Message suited for being used as a context value.
func NewContext(namespace, name string) *Message {
	return &Message{
		Header: map[string]string{
			"namespace": namespace,
			"name":      name,
		},
		Payload: nil,
	}
}

// Creates a Message suited for being used as an event value.
func NewEvent(namespace, name, messageId, dialogRequestId string) *Message {
	m := &Message{
		Header: map[string]string{
			"namespace": namespace,
			"name":      name,
			"messageId": messageId,
		},
		Payload: nil,
	}
	if dialogRequestId != "" {
		m.Header["dialogRequestId"] = dialogRequestId
	}
	return m
}

// Returns a pointer to the underlying Message object.
func (m *Message) GetMessage() *Message {
	return m
}

// Returns the namespace and name as a single string.
func (m *Message) String() string {
	return fmt.Sprintf("%s.%s", m.Header["namespace"], m.Header["name"])
}

// Returns a more specific type for this message.
func (m *Message) Typed() TypedMessage {
	switch m.String() {
	case "Alerts.DeleteAlert":
		return fill(new(DeleteAlert), m)
	case "Alerts.SetAlert":
		return fill(new(SetAlert), m)
	case "AudioPlayer.ClearQueue":
		return fill(new(ClearQueue), m)
	case "AudioPlayer.Play":
		return fill(new(Play), m)
	case "AudioPlayer.Stop":
		return fill(new(Stop), m)
	case "SpeechRecognizer.ExpectSpeech":
		return fill(new(ExpectSpeech), m)
	case "SpeechSynthesizer.Speak":
		return fill(new(Speak), m)
	case "System.Exception":
		return fill(new(Exception), m)
	case "System.SynchronizeState":
		return fill(new(SynchronizeState), m)
	default:
		return m
	}
}

// The AlertEnteredBackground event.
type AlertEnteredBackground struct {
	*Message
	Payload struct {
		Token string `json:"token"`
	} `json:"payload"`
}

func NewAlertEnteredBackground(messageId, token string) *AlertEnteredBackground {
	m := new(AlertEnteredBackground)
	m.Message = NewEvent("Alerts", "AlertEnteredBackground", messageId, "")
	m.Payload.Token = token
	return m
}

// The AlertEnteredForeground event.
type AlertEnteredForeground struct {
	*Message
	Payload struct {
		Token string `json:"token"`
	} `json:"payload"`
}

func NewAlertEnteredForeground(messageId, token string) *AlertEnteredForeground {
	m := new(AlertEnteredForeground)
	m.Message = NewEvent("Alerts", "AlertEnteredForeground", messageId, "")
	m.Payload.Token = token
	return m
}

// The AlertStarted event.
type AlertStarted struct {
	*Message
	Payload struct {
		Token string `json:"token"`
	} `json:"payload"`
}

func NewAlertStarted(messageId, token string) *AlertStarted {
	m := new(AlertStarted)
	m.Message = NewEvent("Alerts", "AlertStarted", messageId, "")
	m.Payload.Token = token
	return m
}

// The AlertStopped event.
type AlertStopped struct {
	*Message
	Payload struct {
		Token string `json:"token"`
	} `json:"payload"`
}

func NewAlertStopped(messageId, token string) *AlertStopped {
	m := new(AlertStopped)
	m.Message = NewEvent("Alerts", "AlertStopped", messageId, "")
	m.Payload.Token = token
	return m
}

// The ClearQueue directive.
type ClearQueue struct {
	*Message
	Payload struct {
		ClearBehavior ClearBehavior `json:"clearBehavior"`
	} `json:"payload"`
}

// The DeleteAlert directive.
type DeleteAlert struct {
	*Message
	Payload struct {
		Token string `json:"token"`
	} `json:"payload"`
}

// The DeleteAlertFailed event.
type DeleteAlertFailed struct {
	*Message
	Payload struct {
		Token string `json:"token"`
	} `json:"payload"`
}

func NewDeleteAlertFailed(messageId, token string) *DeleteAlertFailed {
	m := new(DeleteAlertFailed)
	m.Message = NewEvent("Alerts", "DeleteAlertFailed", messageId, "")
	m.Payload.Token = token
	return m
}

// The DeleteAlertSucceeded event.
type DeleteAlertSucceeded struct {
	*Message
	Payload struct {
		Token string `json:"token"`
	} `json:"payload"`
}

func NewDeleteAlertSucceeded(messageId, token string) *DeleteAlertSucceeded {
	m := new(DeleteAlertSucceeded)
	m.Message = NewEvent("Alerts", "DeleteAlertSucceeded", messageId, "")
	m.Payload.Token = token
	return m
}

// The Exception message.
type Exception struct {
	*Message
	Payload struct {
		Code        string `json:"code"`
		Description string `json:"description"`
	} `json:"payload"`
}

func (m *Exception) Error() string {
	return fmt.Sprintf("%s: %s", m.Payload.Code, m.Payload.Description)
}

// The ExceptionEncountered event.
type ExceptionEncountered struct {
	*Message
	Payload struct {
		UnparsedDirective string `json:"unparsedDirective"`
		Error             struct {
			Type    string `json:"type"`
			Message string `json:"message"`
		} `json:"error"`
	} `json:"payload"`
}

func NewExceptionEncountered(messageId, directive, errorType, errorMessage string) *ExceptionEncountered {
	m := new(ExceptionEncountered)
	m.Message = NewEvent("System", "ExceptionEncountered", messageId, "")
	m.Payload.UnparsedDirective = directive
	m.Payload.Error.Type = errorType
	m.Payload.Error.Message = errorMessage
	return m
}

// The ExpectSpeech directive.
type ExpectSpeech struct {
	*Message
	Payload struct {
		TimeoutInMilliseconds float64 `json:"timeoutInMilliseconds"`
	} `json:"payload"`
}

func (m *ExpectSpeech) Timeout() time.Duration {
	return time.Duration(m.Payload.TimeoutInMilliseconds) * time.Millisecond
}

// The ExpectSpeechTimedOut event.
type ExpectSpeechTimedOut struct {
	*Message
	Payload struct{} `json:"payload"`
}

func NewExpectSpeechTimedOut(messageId string) *ExpectSpeechTimedOut {
	m := new(ExpectSpeechTimedOut)
	m.Message = NewEvent("SpeechRecognizer", "ExpectSpeechTimedOut", messageId, "")
	return m
}

// The Play directive.
type Play struct {
	*Message
	Payload struct {
		AudioItem    AudioItem    `json:"audioItem"`
		PlayBehavior PlayBehavior `json:"playBehavior"`
	} `json:"payload"`
}

func (m *Play) DialogRequestId() string {
	return m.Header["dialogRequestId"]
}

func (m *Play) MessageId() string {
	return m.Header["messageId"]
}

// The PlaybackState context.
type PlaybackState struct {
	*Message
	Payload struct {
		Token                string         `json:"token"`
		OffsetInMilliseconds float64        `json:"offsetInMilliseconds"`
		PlayerActivity       PlayerActivity `json:"playerActivity"`
	} `json:"payload"`
}

func NewPlaybackState(token string, offset time.Duration, activity PlayerActivity) *PlaybackState {
	m := new(PlaybackState)
	m.Message = NewContext("AudioPlayer", "PlaybackState")
	m.Payload.OffsetInMilliseconds = offset.Seconds() * 1000
	m.Payload.PlayerActivity = activity
	m.Payload.Token = token
	return m
}

func (m *PlaybackState) Offset() time.Duration {
	return time.Duration(m.Payload.OffsetInMilliseconds) * time.Millisecond
}

// The Recognize event.
type Recognize struct {
	*Message
	Payload struct {
		Profile string `json:"profile"`
		Format  string `json:"format"`
	} `json:"payload"`
}

func NewRecognize(messageId, dialogRequestId string) *Recognize {
	m := new(Recognize)
	m.Message = NewEvent("SpeechRecognizer", "Recognize", messageId, dialogRequestId)
	m.Payload.Format = "AUDIO_L16_RATE_16000_CHANNELS_1"
	m.Payload.Profile = "CLOSE_TALK"
	return m
}

// The SetAlert directive.
type SetAlert struct {
	*Message
	Payload struct {
		Token         string    `json:"token"`
		Type          AlertType `json:"type"`
		ScheduledTime string    `json:"scheduledTime"`
	} `json:"payload"`
}

// The SetAlertFailed event.
type SetAlertFailed struct {
	*Message
	Payload struct {
		Token string `json:"token"`
	} `json:"payload"`
}

func NewSetAlertFailed(messageId, token string) *SetAlertFailed {
	m := new(SetAlertFailed)
	m.Message = NewEvent("Alerts", "SetAlertFailed", messageId, "")
	m.Payload.Token = token
	return m
}

// The SetAlertSucceeded event.
type SetAlertSucceeded struct {
	*Message
	Payload struct {
		Token string `json:"token"`
	} `json:"payload"`
}

func NewSetAlertSucceeded(messageId, token string) *SetAlertSucceeded {
	m := new(SetAlertSucceeded)
	m.Message = NewEvent("Alerts", "SetAlertSucceeded", messageId, "")
	m.Payload.Token = token
	return m
}

// The Speak directive.
type Speak struct {
	*Message
	Payload struct {
		Format string
		URL    string
	} `json:"payload"`
}

func (m *Speak) ContentId() string {
	if !strings.HasPrefix(m.Payload.URL, "cid:") {
		return ""
	}
	return m.Payload.URL[4:]
}

// The ResetUserInactivity directive.
type ResetUserInactivity struct {
	*Message
	Payload struct{} `json:"payload"`
}

// The Stop directive.
type Stop struct {
	*Message
	Payload struct{} `json:"payload"`
}

// The SynchronizeState event.
type SynchronizeState struct {
	*Message
	Payload struct{} `json:"payload"`
}

func NewSynchronizeState(messageId string) *SynchronizeState {
	m := new(SynchronizeState)
	m.Message = NewEvent("System", "SynchronizeState", messageId, "")
	return m
}

// The UserInactivityReport event.
type UserInactivityReport struct {
	*Message
	Payload struct {
		InactiveTimeInSeconds float64 `json:"inactiveTimeInSeconds"`
	} `json:"payload"`
}

func NewUserInactivityReport(messageId string, inactiveTime time.Duration) *UserInactivityReport {
	m := new(UserInactivityReport)
	m.Message = NewEvent("System", "UserInactivityReport", messageId, "")
	m.Payload.InactiveTimeInSeconds = inactiveTime.Seconds()
	return m
}

// Convenience method to set up an empty typed message object from a raw Message.
func fill(dst TypedMessage, src *Message) TypedMessage {
	v := reflect.ValueOf(dst).Elem()
	v.FieldByName("Message").Set(reflect.ValueOf(src))
	payload := v.FieldByName("Payload")
	if payload.Kind() != reflect.Struct {
		return dst
	}
	json.Unmarshal(src.Payload, payload.Addr().Interface())
	return dst
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/binary"
	"fmt"
	"strconv"
	"io"
	"net"
	"github.com/lbarman/prifi/dcnet"
	//log2 "github.com/lbarman/prifi/log"
)

// Number of bytes of cell payload to reserve for connection header, length
const socksHeaderLength = 6

func startClient(clientId int, relayHostAddr string, nClients int, nTrustees int, payloadLength int, useSocksProxy bool) {
	fmt.Printf("startClient %d\n", clientId)

	//crypto parameters
	tg := dcnet.TestSetup(nil, suite, factory, nClients, nTrustees)
	me := tg.Clients[clientId]
	clientPayloadSize := me.Coder.ClientCellSize(payloadLength)

	relayConn := connectToRelay(relayHostAddr, clientId)

	//initiate downstream stream
	dataFromRelay := make(chan dataWithConnectionId)
	go readDataFromRelay(relayConn, dataFromRelay)

	println("client", clientId, "connected")

	// We're the "slot owner" - start a socks relay
	socksProxyNewConnections    := make(chan net.Conn)
	socksProxyData              := make(chan []byte)
	socksProxyConnClosed        := make(chan int)
	socksProxyActiveConnections := make([]net.Conn, 1) // reserve socksProxyActiveConnections[0]

	if(useSocksProxy){
		port := ":" + strconv.Itoa(1080+clientId)
		go startSocksProxy(port, socksProxyNewConnections)
	}

	// This will hold the data to be sent later on to the relay, anonymized
	dataForRelayBuffer := make([][]byte, 0)

	// Client/proxy main loop
	totupcells := uint64(0)
	totupbytes := uint64(0)
	for {
		select {

			// New TCP connection to the SOCKS proxy
			case conn := <-socksProxyNewConnections:
				newClientId := len(socksProxyActiveConnections)
				socksProxyActiveConnections = append(socksProxyActiveConnections, conn)
				go readDataFromSocksProxy(newClientId, payloadLength, conn, socksProxyData, socksProxyConnClosed)

			// Data to anonymize from SOCKS proxy
			case data := <-socksProxyData:
				dataForRelayBuffer = append(dataForRelayBuffer, data)

			//connection closed from SOCKS proxy
			case clientId := <-socksProxyConnClosed:
				socksProxyActiveConnections[clientId] = nil

			//downstream slice from relay (normal DC-net cycle)
			case dataWithConnId := <-dataFromRelay:
				print(".")

				connId := dataWithConnId.connectionId

				//Handle the connections, forwards the downstream slice to the SOCKS proxy
				if connId > 0 && connId < len(socksProxyActiveConnections) && socksProxyActiveConnections[connId] != nil {
					data       := dataWithConnId.data
					dataLength := len(data)

					if dataLength > 0 {

						//if there is no socks proxy, nothing to do (useless case indeed, only for debug)
						if useSocksProxy {
							n, err := socksProxyActiveConnections[clientId].Write(data)
							if n < dataLength {
								panic("Write to socks proxy: expected "+strconv.Itoa(dataLength)+" bytes, got "+strconv.Itoa(n)+", " + err.Error())
							}
						}
					} else {
						// Relay indicating EOF on this conn
						fmt.Printf("Relay to client : closed conn %d", connId)
						socksProxyActiveConnections[clientId].Close()
					}
				}

				// Should account the downstream cell in the history

				// Produce and ship the next upstream slice
				writeNextUpstreamSlice(dataForRelayBuffer, payloadLength, clientPayloadSize, relayConn, me)

				//statistics
				totupcells++
				totupbytes += uint64(payloadLength)
				//fmt.Printf("sent %d upstream cells, %d bytes\n", totupcells, totupbytes)
			}
	}
}

/*
 * Creates the next cell
 */

func writeNextUpstreamSlice(dataForRelayBuffer [][]byte, payloadLength int, clientPayloadSize int, relayConn net.Conn, me *dcnet.TestNode) {
	var nextUpstreamBytes []byte
	if len(dataForRelayBuffer) > 0 {
		nextUpstreamBytes  = dataForRelayBuffer[0]
		dataForRelayBuffer = dataForRelayBuffer[1:]
		//fmt.Printf("\n^ %v (len : %d)\n", p)
	}

	//produce the next upstream cell
	upstreamSlice := me.Coder.ClientEncode(nextUpstreamBytes, payloadLength, me.History)

	if len(upstreamSlice) != clientPayloadSize {
		panic("Client slice wrong size, expected "+strconv.Itoa(clientPayloadSize)+", but got "+strconv.Itoa(len(upstreamSlice)))
	}

	n, err := relayConn.Write(upstreamSlice)
	if n != len(upstreamSlice) {
		panic("Client write to relay error, expected writing "+strconv.Itoa(len(upstreamSlice))+", but wrote "+strconv.Itoa(n)+", err : " + err.Error())
	}
}

/*
 * RELAY CONNECTION
 */

func connectToRelay(relayHost string, connectionId int) net.Conn {
	conn, err := net.Dial("tcp", relayHost)
	if err != nil {
		panic("Can't connect to relay:" + err.Error())
	}

	// Tell the relay our client or trustee number
	b := make([]byte, 1)
	b[0] = byte(connectionId)
	n, err := conn.Write(b)

	if n < 1 || err != nil {
		panic("Error writing to socket:" + err.Error())
	}

	return conn
}

func readDataFromRelay(relayConn net.Conn, datadataFromRelay chan<- dataWithConnectionId) {
	header := [6]byte{}
	totcells := uint64(0)
	totbytes := uint64(0)

	for {
		// Read the next (downstream) header from the relay
		n, err := io.ReadFull(relayConn, header[:])

		if n != len(header) {
			panic("clientReadRelay: " + err.Error())
		}

		connectionId := int(binary.BigEndian.Uint32(header[0:4]))
		dataLength := int(binary.BigEndian.Uint16(header[4:6]))

		// Read the downstream data
		data := make([]byte, dataLength)
		n, err = io.ReadFull(relayConn, data)

		if n != dataLength {
			panic("readDataFromRelay: read data length ("+strconv.Itoa(n)+") not matching expected length ("+strconv.Itoa(dataLength)+")" + err.Error())
		}

		datadataFromRelay <- dataWithConnectionId{connectionId, data}

		totcells++
		totbytes += uint64(dataLength)
	}
}

/*
 * SOCKS PROXY
 */

func startSocksProxy(port string, newConnections chan<- net.Conn) {
	fmt.Printf("Listening on port %s\n", port)

	lsock, err := net.Listen("tcp", port)

	if err != nil {
		fmt.Printf("Can't open listen socket at port %s: %s", port, err.Error())
		return
	}

	for {
		conn, err := lsock.Accept()
		fmt.Printf("Accept on port %s\n", port)

		if err != nil {
			lsock.Close()
			return
		}
		newConnections <- conn
	}
}

func readDataFromSocksProxy(clientId int, payloadLength int, conn net.Conn, data chan<- []byte, closed chan<- int) {
	for {
		// Read up to a cell worth of data to send upstream
		buffer := make([]byte, payloadLength)
		n, err := conn.Read(buffer[socksHeaderLength:])

		// Encode the connection number and actual data length
		binary.BigEndian.PutUint32(buffer[0:4], uint32(clientId))
		binary.BigEndian.PutUint16(buffer[4:6], uint16(n))

		data <- buffer

		// Connection error or EOF?
		if n == 0 {
			if err == io.EOF {
				println("clientUpload: EOF, closing")
			} else {
				println("clientUpload: " + err.Error())
			}
			conn.Close()
			closed <- clientId // signal that channel is closed
			return
		}
	}
}<commit_msg>Adds keys struct to client<commit_after>package main

import (
	"encoding/binary"
	"fmt"
	"strconv"
	"io"
	"net"
	"github.com/lbarman/prifi/dcnet"
	//log2 "github.com/lbarman/prifi/log"
)

// Number of bytes of cell payload to reserve for connection header, length
const socksHeaderLength = 6

type ClientParams struct {

	// Asymmetric keypair for this client
	pub abstract.Point
	pri abstract.Secret

	trusteekeys      []abstract.Point  // each trustee's public key
	sharedsecrets []abstract.Cipher // shared secrets
}

func startClient(clientId int, relayHostAddr string, nClients int, nTrustees int, payloadLength int, useSocksProxy bool) {
	fmt.Printf("startClient %d\n", clientId)

	//crypto parameters
	tg := dcnet.TestSetup(nil, suite, factory, nClients, nTrustees)
	me := tg.Clients[clientId]
	clientPayloadSize := me.Coder.ClientCellSize(payloadLength)

	relayConn := connectToRelay(relayHostAddr, clientId)

	//initiate downstream stream
	dataFromRelay := make(chan dataWithConnectionId)
	go readDataFromRelay(relayConn, dataFromRelay)

	println("client", clientId, "connected")

	// We're the "slot owner" - start a socks relay
	socksProxyNewConnections    := make(chan net.Conn)
	socksProxyData              := make(chan []byte)
	socksProxyConnClosed        := make(chan int)
	socksProxyActiveConnections := make([]net.Conn, 1) // reserve socksProxyActiveConnections[0]

	if(useSocksProxy){
		port := ":" + strconv.Itoa(1080+clientId)
		go startSocksProxy(port, socksProxyNewConnections)
	}

	// This will hold the data to be sent later on to the relay, anonymized
	dataForRelayBuffer := make([][]byte, 0)

	// Client/proxy main loop
	totupcells := uint64(0)
	totupbytes := uint64(0)
	for {
		select {

			// New TCP connection to the SOCKS proxy
			case conn := <-socksProxyNewConnections:
				newClientId := len(socksProxyActiveConnections)
				socksProxyActiveConnections = append(socksProxyActiveConnections, conn)
				go readDataFromSocksProxy(newClientId, payloadLength, conn, socksProxyData, socksProxyConnClosed)

			// Data to anonymize from SOCKS proxy
			case data := <-socksProxyData:
				dataForRelayBuffer = append(dataForRelayBuffer, data)

			//connection closed from SOCKS proxy
			case clientId := <-socksProxyConnClosed:
				socksProxyActiveConnections[clientId] = nil

			//downstream slice from relay (normal DC-net cycle)
			case dataWithConnId := <-dataFromRelay:
				print(".")

				connId := dataWithConnId.connectionId

				//Handle the connections, forwards the downstream slice to the SOCKS proxy
				if connId > 0 && connId < len(socksProxyActiveConnections) && socksProxyActiveConnections[connId] != nil {
					data       := dataWithConnId.data
					dataLength := len(data)

					if dataLength > 0 {

						//if there is no socks proxy, nothing to do (useless case indeed, only for debug)
						if useSocksProxy {
							n, err := socksProxyActiveConnections[clientId].Write(data)
							if n < dataLength {
								panic("Write to socks proxy: expected "+strconv.Itoa(dataLength)+" bytes, got "+strconv.Itoa(n)+", " + err.Error())
							}
						}
					} else {
						// Relay indicating EOF on this conn
						fmt.Printf("Relay to client : closed conn %d", connId)
						socksProxyActiveConnections[clientId].Close()
					}
				}

				// Should account the downstream cell in the history

				// Produce and ship the next upstream slice
				writeNextUpstreamSlice(dataForRelayBuffer, payloadLength, clientPayloadSize, relayConn, me)

				//statistics
				totupcells++
				totupbytes += uint64(payloadLength)
				//fmt.Printf("sent %d upstream cells, %d bytes\n", totupcells, totupbytes)
			}
	}
}

/*
 * Creates the next cell
 */

func writeNextUpstreamSlice(dataForRelayBuffer [][]byte, payloadLength int, clientPayloadSize int, relayConn net.Conn, me *dcnet.TestNode) {
	var nextUpstreamBytes []byte
	if len(dataForRelayBuffer) > 0 {
		nextUpstreamBytes  = dataForRelayBuffer[0]
		dataForRelayBuffer = dataForRelayBuffer[1:]
		//fmt.Printf("\n^ %v (len : %d)\n", p)
	}

	//produce the next upstream cell
	upstreamSlice := me.Coder.ClientEncode(nextUpstreamBytes, payloadLength, me.History)

	if len(upstreamSlice) != clientPayloadSize {
		panic("Client slice wrong size, expected "+strconv.Itoa(clientPayloadSize)+", but got "+strconv.Itoa(len(upstreamSlice)))
	}

	n, err := relayConn.Write(upstreamSlice)
	if n != len(upstreamSlice) {
		panic("Client write to relay error, expected writing "+strconv.Itoa(len(upstreamSlice))+", but wrote "+strconv.Itoa(n)+", err : " + err.Error())
	}
}

/*
 * RELAY CONNECTION
 */

func connectToRelay(relayHost string, connectionId int) net.Conn {
	conn, err := net.Dial("tcp", relayHost)
	if err != nil {
		panic("Can't connect to relay:" + err.Error())
	}

	// Tell the relay our client or trustee number
	b := make([]byte, 1)
	b[0] = byte(connectionId)
	n, err := conn.Write(b)

	if n < 1 || err != nil {
		panic("Error writing to socket:" + err.Error())
	}

	return conn
}

func readDataFromRelay(relayConn net.Conn, datadataFromRelay chan<- dataWithConnectionId) {
	header := [6]byte{}
	totcells := uint64(0)
	totbytes := uint64(0)

	for {
		// Read the next (downstream) header from the relay
		n, err := io.ReadFull(relayConn, header[:])

		if n != len(header) {
			panic("clientReadRelay: " + err.Error())
		}

		connectionId := int(binary.BigEndian.Uint32(header[0:4]))
		dataLength := int(binary.BigEndian.Uint16(header[4:6]))

		// Read the downstream data
		data := make([]byte, dataLength)
		n, err = io.ReadFull(relayConn, data)

		if n != dataLength {
			panic("readDataFromRelay: read data length ("+strconv.Itoa(n)+") not matching expected length ("+strconv.Itoa(dataLength)+")" + err.Error())
		}

		datadataFromRelay <- dataWithConnectionId{connectionId, data}

		totcells++
		totbytes += uint64(dataLength)
	}
}

/*
 * SOCKS PROXY
 */

func startSocksProxy(port string, newConnections chan<- net.Conn) {
	fmt.Printf("Listening on port %s\n", port)

	lsock, err := net.Listen("tcp", port)

	if err != nil {
		fmt.Printf("Can't open listen socket at port %s: %s", port, err.Error())
		return
	}

	for {
		conn, err := lsock.Accept()
		fmt.Printf("Accept on port %s\n", port)

		if err != nil {
			lsock.Close()
			return
		}
		newConnections <- conn
	}
}

func readDataFromSocksProxy(clientId int, payloadLength int, conn net.Conn, data chan<- []byte, closed chan<- int) {
	for {
		// Read up to a cell worth of data to send upstream
		buffer := make([]byte, payloadLength)
		n, err := conn.Read(buffer[socksHeaderLength:])

		// Encode the connection number and actual data length
		binary.BigEndian.PutUint32(buffer[0:4], uint32(clientId))
		binary.BigEndian.PutUint16(buffer[4:6], uint16(n))

		data <- buffer

		// Connection error or EOF?
		if n == 0 {
			if err == io.EOF {
				println("clientUpload: EOF, closing")
			} else {
				println("clientUpload: " + err.Error())
			}
			conn.Close()
			closed <- clientId // signal that channel is closed
			return
		}
	}
}<|endoftext|>


<commit_before>package libvirt

import (
	"fmt"
	"log"
	"strings"
	"time"

	"github.com/mitchellh/multistep"
	"github.com/mitchellh/packer/packer"
	"gopkg.in/alexzorin/libvirt-go.v2"
)

const KeyLeftShift uint32 = 0xFFE1

type bootCommandTemplateData struct {
	HTTPIP   string
	HTTPPort uint
	Name     string
}

// This step "types" the boot command into the VM over VNC.
//
// Uses:
//   config *config
//   http_port int
//   ui     packer.Ui
//
// Produces:
//   <nothing>
type stepTypeBootCommand struct{}

func (s *stepTypeBootCommand) Run(state multistep.StateBag) multistep.StepAction {
	config := state.Get("config").(*Config)
	//	httpPort := state.Get("http_port").(uint)
	//	hostIp := state.Get("host_ip").(string)
	ui := state.Get("ui").(packer.Ui)

	var lvd libvirt.VirDomain
	lv, err := libvirt.NewVirConnection(config.LibvirtUrl)
	if err != nil {
		err := fmt.Errorf("Error connecting to libvirt: %s", err)
		state.Put("error", err)
		ui.Error(err.Error())
		return multistep.ActionHalt
	}
	defer lv.CloseConnection()
	if lvd, err = lv.LookupDomainByName(config.VMName); err != nil {
		err := fmt.Errorf("Error lookup domain: %s", err)
		state.Put("error", err)
		ui.Error(err.Error())
		return multistep.ActionHalt
	}
	defer lvd.Free()

	//	tplData := &bootCommandTemplateData{
	//		hostIp,
	//		httpPort,
	//		config.VMName,
	//	}

	ui.Say("Typing the boot command...")
	for _, command := range config.BootCommand {
		//		command, err := config.tpl.Process(command, tplData)
		//		if err != nil {
		//			err := fmt.Errorf("Error preparing boot command: %s", err)
		//			state.Put("error", err)
		//			ui.Error(err.Error())
		//			return multistep.ActionHalt
		//		}

		// Check for interrupts between typing things so we can cancel
		// since this isn't the fastest thing.
		if _, ok := state.GetOk(multistep.StateCancelled); ok {
			return multistep.ActionHalt
		}

		sendBootString(lvd, command)
	}

	return multistep.ActionContinue
}

func (*stepTypeBootCommand) Cleanup(multistep.StateBag) {}

func sendBootString(d libvirt.VirDomain, original string) {
	//	shiftedChars := "~!@#$%^&*()_+{}|:\"<>?"
	var err error
	var ok bool
	var key uint

	for len(original) > 0 {
		//		var keyCode uint
		//		keyShift := false

		if strings.HasPrefix(original, "<wait>") {
			log.Printf("Special code '<wait>' found, sleeping one second")
			time.Sleep(1 * time.Second)
			original = original[len("<wait>"):]
			continue
		}

		if strings.HasPrefix(original, "<wait5>") {
			log.Printf("Special code '<wait5>' found, sleeping 5 seconds")
			time.Sleep(5 * time.Second)
			original = original[len("<wait5>"):]
			continue
		}

		if strings.HasPrefix(original, "<wait10>") {
			log.Printf("Special code '<wait10>' found, sleeping 10 seconds")
			time.Sleep(10 * time.Second)
			original = original[len("<wait10>"):]
			continue
		}

		if strings.HasPrefix(original, "<esc>") {
			d.SendKey(libvirt.VIR_KEYCODE_SET_RFB, 500, []uint{ecodes["<esc>"]}, 0)
			original = original[len("<esc>"):]
		}
		if strings.HasPrefix(original, "<enter>") {
			d.SendKey(libvirt.VIR_KEYCODE_SET_RFB, 500, []uint{ecodes["<enter>"]}, 0)
			original = original[len("<enter>"):]
		}

		char := original[0]
		log.Printf("try to find code for char %s", string(char))
		if key, ok = ecodes[string(char)]; ok {
			log.Printf("find code for char %s %d", string(char), key)
			//			keyShift = unicode.IsUpper(r) || strings.ContainsRune(shiftedChars, r)
		} else {
			log.Printf("can't find code for char %s", string(char))
			continue
		}
		original = original[1:]
		//VIR_KEYCODE_SET_LINUX, VIR_KEYCODE_SET_USB, VIR_KEYCODE_SET_RFB, VIR_KEYCODE_SET_WIN32, VIR_KEYCODE_SET_XT_KBD
		log.Printf("send code %d", key)
		if err = d.SendKey(libvirt.VIR_KEYCODE_SET_RFB, 500, []uint{key}, 0); err != nil {
			log.Printf("Sending code %d failed: %s", key, err.Error())
		}
	}

}
<commit_msg>fix<commit_after>package libvirt

import (
	"fmt"
	"log"
	"strings"
	"time"
	"unicode/utf8"

	"github.com/mitchellh/multistep"
	"github.com/mitchellh/packer/packer"
	"gopkg.in/alexzorin/libvirt-go.v2"
)

const KeyLeftShift uint32 = 0xFFE1

type bootCommandTemplateData struct {
	HTTPIP   string
	HTTPPort uint
	Name     string
}

// This step "types" the boot command into the VM over VNC.
//
// Uses:
//   config *config
//   http_port int
//   ui     packer.Ui
//
// Produces:
//   <nothing>
type stepTypeBootCommand struct{}

func (s *stepTypeBootCommand) Run(state multistep.StateBag) multistep.StepAction {
	config := state.Get("config").(*Config)
	//	httpPort := state.Get("http_port").(uint)
	//	hostIp := state.Get("host_ip").(string)
	ui := state.Get("ui").(packer.Ui)

	var lvd libvirt.VirDomain
	lv, err := libvirt.NewVirConnection(config.LibvirtUrl)
	if err != nil {
		err := fmt.Errorf("Error connecting to libvirt: %s", err)
		state.Put("error", err)
		ui.Error(err.Error())
		return multistep.ActionHalt
	}
	defer lv.CloseConnection()
	if lvd, err = lv.LookupDomainByName(config.VMName); err != nil {
		err := fmt.Errorf("Error lookup domain: %s", err)
		state.Put("error", err)
		ui.Error(err.Error())
		return multistep.ActionHalt
	}
	defer lvd.Free()

	//	tplData := &bootCommandTemplateData{
	//		hostIp,
	//		httpPort,
	//		config.VMName,
	//	}

	ui.Say("Typing the boot command...")
	for _, command := range config.BootCommand {
		//		command, err := config.tpl.Process(command, tplData)
		//		if err != nil {
		//			err := fmt.Errorf("Error preparing boot command: %s", err)
		//			state.Put("error", err)
		//			ui.Error(err.Error())
		//			return multistep.ActionHalt
		//		}

		// Check for interrupts between typing things so we can cancel
		// since this isn't the fastest thing.
		if _, ok := state.GetOk(multistep.StateCancelled); ok {
			return multistep.ActionHalt
		}

		sendBootString(lvd, command)
	}

	return multistep.ActionContinue
}

func (*stepTypeBootCommand) Cleanup(multistep.StateBag) {}

func sendBootString(d libvirt.VirDomain, original string) {
	//	shiftedChars := "~!@#$%^&*()_+{}|:\"<>?"
	var err error
	var ok bool
	var key uint

	for len(original) > 0 {
		//		var keyCode uint
		//		keyShift := false

		if strings.HasPrefix(original, "<wait>") {
			log.Printf("Special code '<wait>' found, sleeping one second")
			time.Sleep(1 * time.Second)
			original = original[len("<wait>"):]
			continue
		}

		if strings.HasPrefix(original, "<wait5>") {
			log.Printf("Special code '<wait5>' found, sleeping 5 seconds")
			time.Sleep(5 * time.Second)
			original = original[len("<wait5>"):]
			continue
		}

		if strings.HasPrefix(original, "<wait10>") {
			log.Printf("Special code '<wait10>' found, sleeping 10 seconds")
			time.Sleep(10 * time.Second)
			original = original[len("<wait10>"):]
			continue
		}

		if strings.HasPrefix(original, "<esc>") {
			d.SendKey(libvirt.VIR_KEYCODE_SET_RFB, 400, []uint{ecodes["<esc>"]}, 0)
			original = original[len("<esc>"):]
		}
		if strings.HasPrefix(original, "<enter>") {
			d.SendKey(libvirt.VIR_KEYCODE_SET_RFB, 400, []uint{ecodes["<enter>"]}, 0)
			original = original[len("<enter>"):]
		}

		r, size := utf8.DecodeRuneInString(original)
		original = original[size:]
		key = ecodes[string(r)]
		log.Printf("find code for char %s %d", string(r), key)
		//VIR_KEYCODE_SET_LINUX, VIR_KEYCODE_SET_USB, VIR_KEYCODE_SET_RFB, VIR_KEYCODE_SET_WIN32, VIR_KEYCODE_SET_XT_KBD
		log.Printf("send code %d", key)
		if err = d.SendKey(libvirt.VIR_KEYCODE_SET_RFB, 400, []uint{key}, 0); err != nil {
			log.Printf("Sending code %d failed: %s", key, err.Error())
		}
	}

}
<|endoftext|>