<commit_before>package main

import (
	"fmt"
	"os"
	"time"

	"github.com/StreamBoat/kodi_jsonrpc"
	"github.com/pdf/kodi-callback-daemon/config"
	"github.com/pdf/kodi-callback-daemon/hyperion"
	"github.com/pdf/kodi-callback-daemon/kodi"
	"github.com/pdf/kodi-callback-daemon/shell"

	log "github.com/Sirupsen/logrus"
)

const (
	VERSION = "1.0.0"
)

var (
	cfg config.Config
	k   kodi_jsonrpc.Connection
)

// usage simply prints the invocation requirements.
func usage() {
	fmt.Fprintf(os.Stderr, "\nKodi Callback Daemon v%s\n\nUsage: %s [configFile]\n\n", VERSION, os.Args[0])
	os.Exit(1)
}

// init ensures we have a config path argument, and loads the configuration.
func init() {
	if len(os.Args) < 2 {
		usage()
	}

	// Initialize logger
	cfg = config.Load(os.Args[1])
	if cfg.Debug != nil && *cfg.Debug == true {
		log.SetLevel(log.DebugLevel)
	} else {
		log.SetLevel(log.InfoLevel)
	}
}

// execute iterates through a list of callbacks, and sends them to the backend
// defined in the `backend` property.
func execute(callbacks []interface{}) {
	for i := range callbacks {
		m := callbacks[i].(map[string]interface{})

		switch m[`backend`] {
		case `hyperion`:
			if cfg.Hyperion != nil {
				hyperion.Execute(m)
			}

		case `kodi`, `xbmc`:
			kodi.Execute(&k, m)

		case `shell`:
			shell.Execute(m)

		default:
			log.WithField(`backend`, m[`backend`]).Warn(`Unknown backend`)
		}
	}
}

// callbacksByType takes a type to match, and a list of callbacks. If the
// callback has a `types` property, and that contains a matching type, the
// callback is added to the returned list.  A callback without a `types`
// property will always be returned.
func callbacksByType(matchType string, callbacks []interface{}) []interface{} {
	result := make([]interface{}, 0)
	var cb map[string]interface{}

	for i := range callbacks {
		// Access internal callback map.
		cb = callbacks[i].(map[string]interface{})

		switch cb[`types`].(type) {
		// We have a list of types.
		case []interface{}:
			// Access internal types slice.
			cbTypes, ok := cb[`types`].([]interface{})
			if ok == false {
				log.Fatal(`Couldn't understand 'types' array, check your configuration.`)
			}
			for j := range cbTypes {
				if cbTypes[j].(string) == matchType {
					// Matched the required type, add this callback to the results.
					result = append(result, cb)
				}
			}
		// If there is no valid `types` property, add this callback to the results.
		default:
			result = append(result, cb)
		}
	}

	return result
}

// main program loop.
func main() {
	var host config.Host
	// Connect to Kodi, this is required.
	kodi_timeout := time.Duration(0)
	if cfg.Kodi != nil {
		if cfg.Kodi.Timeout != nil {
			kodi_timeout = *cfg.Kodi.Timeout
		}
		host = *cfg.Kodi
	} else if cfg.XBMC != nil {
		if cfg.XBMC.Timeout != nil {
			kodi_timeout = *cfg.XBMC.Timeout
		}
		host = *cfg.XBMC
	} else {
		log.Fatal(`You must provide Kodi/XBMC connection details in your configuration`)
	}
	k, err := kodi_jsonrpc.New(
		fmt.Sprintf(`%s:%d`, host.Address, host.Port),
		kodi_timeout,
	)

	defer k.Close()
	if err != nil {
		log.WithField(`error`, err).Fatal(`Failed to obtain Kodi connection`)
	}

	// If the configuration specifies a Hyperion connection, use it.
	if cfg.Hyperion != nil {
		hyperion.Connect(fmt.Sprintf(`%s:%d`, cfg.Hyperion.Address, cfg.Hyperion.Port))
		defer hyperion.Close()
	}

	// Get callbacks from configuration.
	callbacks := cfg.Callbacks.(map[string]interface{})

	// Execute callbacks for the special `Startup` notification.
	if callbacks[`Startup`] != nil {
		execute(callbacks[`Startup`].([]interface{}))
	}

	// Loop while reading from Kodi.
	for {
		// Read from Kodi.
		notification := <-k.Notifications

		// Match Kodi notification to our configured callbacks.
		if callbacks[notification.Method] != nil {
			cbs := callbacks[notification.Method].([]interface{})
			// The Player.OnPlay notification supports an filtering by item type.
			if notification.Method == `Player.OnPlay` {
				cbs = callbacksByType(notification.Params.Data.Item.Type, cbs)
			}
			execute(cbs)
		}
	}
}
<commit_msg>Bump version to v1.0.1<commit_after>package main

import (
	"fmt"
	"os"
	"time"

	"github.com/StreamBoat/kodi_jsonrpc"
	"github.com/pdf/kodi-callback-daemon/config"
	"github.com/pdf/kodi-callback-daemon/hyperion"
	"github.com/pdf/kodi-callback-daemon/kodi"
	"github.com/pdf/kodi-callback-daemon/shell"

	log "github.com/Sirupsen/logrus"
)

const (
	VERSION = "1.0.1"
)

var (
	cfg config.Config
	k   kodi_jsonrpc.Connection
)

// usage simply prints the invocation requirements.
func usage() {
	fmt.Fprintf(os.Stderr, "\nKodi Callback Daemon v%s\n\nUsage: %s [configFile]\n\n", VERSION, os.Args[0])
	os.Exit(1)
}

// init ensures we have a config path argument, and loads the configuration.
func init() {
	if len(os.Args) < 2 {
		usage()
	}

	// Initialize logger
	cfg = config.Load(os.Args[1])
	if cfg.Debug != nil && *cfg.Debug == true {
		log.SetLevel(log.DebugLevel)
	} else {
		log.SetLevel(log.InfoLevel)
	}
}

// execute iterates through a list of callbacks, and sends them to the backend
// defined in the `backend` property.
func execute(callbacks []interface{}) {
	for i := range callbacks {
		m := callbacks[i].(map[string]interface{})

		switch m[`backend`] {
		case `hyperion`:
			if cfg.Hyperion != nil {
				hyperion.Execute(m)
			}

		case `kodi`, `xbmc`:
			kodi.Execute(&k, m)

		case `shell`:
			shell.Execute(m)

		default:
			log.WithField(`backend`, m[`backend`]).Warn(`Unknown backend`)
		}
	}
}

// callbacksByType takes a type to match, and a list of callbacks. If the
// callback has a `types` property, and that contains a matching type, the
// callback is added to the returned list.  A callback without a `types`
// property will always be returned.
func callbacksByType(matchType string, callbacks []interface{}) []interface{} {
	result := make([]interface{}, 0)
	var cb map[string]interface{}

	for i := range callbacks {
		// Access internal callback map.
		cb = callbacks[i].(map[string]interface{})

		switch cb[`types`].(type) {
		// We have a list of types.
		case []interface{}:
			// Access internal types slice.
			cbTypes, ok := cb[`types`].([]interface{})
			if ok == false {
				log.Fatal(`Couldn't understand 'types' array, check your configuration.`)
			}
			for j := range cbTypes {
				if cbTypes[j].(string) == matchType {
					// Matched the required type, add this callback to the results.
					result = append(result, cb)
				}
			}
		// If there is no valid `types` property, add this callback to the results.
		default:
			result = append(result, cb)
		}
	}

	return result
}

// main program loop.
func main() {
	var host config.Host
	// Connect to Kodi, this is required.
	kodi_timeout := time.Duration(0)
	if cfg.Kodi != nil {
		if cfg.Kodi.Timeout != nil {
			kodi_timeout = *cfg.Kodi.Timeout
		}
		host = *cfg.Kodi
	} else if cfg.XBMC != nil {
		if cfg.XBMC.Timeout != nil {
			kodi_timeout = *cfg.XBMC.Timeout
		}
		host = *cfg.XBMC
	} else {
		log.Fatal(`You must provide Kodi/XBMC connection details in your configuration`)
	}
	k, err := kodi_jsonrpc.New(
		fmt.Sprintf(`%s:%d`, host.Address, host.Port),
		kodi_timeout,
	)

	defer k.Close()
	if err != nil {
		log.WithField(`error`, err).Fatal(`Failed to obtain Kodi connection`)
	}

	// If the configuration specifies a Hyperion connection, use it.
	if cfg.Hyperion != nil {
		hyperion.Connect(fmt.Sprintf(`%s:%d`, cfg.Hyperion.Address, cfg.Hyperion.Port))
		defer hyperion.Close()
	}

	// Get callbacks from configuration.
	callbacks := cfg.Callbacks.(map[string]interface{})

	// Execute callbacks for the special `Startup` notification.
	if callbacks[`Startup`] != nil {
		execute(callbacks[`Startup`].([]interface{}))
	}

	// Loop while reading from Kodi.
	for {
		// Read from Kodi.
		notification := <-k.Notifications

		// Match Kodi notification to our configured callbacks.
		if callbacks[notification.Method] != nil {
			cbs := callbacks[notification.Method].([]interface{})
			// The Player.OnPlay notification supports an filtering by item type.
			if notification.Method == `Player.OnPlay` {
				cbs = callbacksByType(notification.Params.Data.Item.Type, cbs)
			}
			execute(cbs)
		}
	}
}
<|endoftext|>


<commit_before>package qb

import (
	"fmt"
	"strings"
)

// NewTable generates a new table pointer given table name, column and table constraints
func NewTable(driver string, name string, columns []Column, constraints []Constraint) *Table {
	return &Table{
		name:        name,
		columns:     columns,
		constraints: constraints,
		primaryCols: []string{},
		refs:        []ref{},
		driver:      driver,
	}
}

// Table is the base abstraction for any sql table
type Table struct {
	name        string
	columns     []Column
	constraints []Constraint
	primaryCols []string
	refs        []ref
	driver      string
}

// Name returns the table name
func (t *Table) Name() string {
	return t.name
}

// SQL generates create table syntax of table
func (t *Table) SQL() string {

	dialect := NewDialect(t.driver)

	cols := []string{}
	for _, v := range t.columns {
		cols = append(cols, v.SQL(t.driver))
	}

	constraints := []string{}

	// build primary key constraints using primaryCols
	if len(t.primaryCols) > 0 {
		constraints = append(constraints, fmt.Sprintf("PRIMARY KEY (%s)", strings.Join(t.primaryCols, ", ")))
	}

	// build foreign key constraints using refCols
	for _, ref := range t.refs {
		constraints = append(constraints, fmt.Sprintf("FOREIGN KEY (%s) REFERENCES %s(%s)", strings.Join(ref.cols, ", "), ref.refTable, strings.Join(ref.refCols, ", ")))
	}

	for _, v := range t.constraints {
		constraints = append(constraints, v.Name)
	}

	query := dialect.CreateTable(t.name, cols, constraints).Query()

	return query.SQL()
}

// AddColumn appends a new column to current table
func (t *Table) AddColumn(column Column) {
	t.columns = append(t.columns, column)
}

// AddConstraint appends a new constraint to current table
func (t *Table) AddConstraint(c Constraint) {
	t.constraints = append(t.constraints, c)
}

// AddPrimary appends a primary column that will be lazily built as a primary key constraint
func (t *Table) AddPrimary(col string) {
	t.primaryCols = append(t.primaryCols, col)
}

type ref struct {
	cols     []string
	refTable string
	refCols  []string
}

// AddRef appends a new reference struct that will be lazily built as a foreign key constraint
func (t *Table) AddRef(col string, refTable string, refCol string) {

	if len(t.refs) > 0 {
		for k, ref := range t.refs {
			if refTable == ref.refTable {
				t.refs[k].cols = append(t.refs[k].cols, fmt.Sprintf("%s", col))
				t.refs[k].refCols = append(t.refs[k].refCols, fmt.Sprintf("%s", refCol))
				return
			}
		}
	}

	r := ref{[]string{}, refTable, []string{}}
	r.cols = append(r.cols, col)
	r.refCols = append(r.refCols, refCol)
	t.refs = append(t.refs, r)
}

// Constraints returns the constraint slice of current table
func (t *Table) Constraints() []Constraint {
	return t.constraints
}

// Insert creates an insert statement for the table name
func (t *Table) Insert(kv map[string]interface{}) *Dialect {

	keys := []string{}
	values := []interface{}{}

	for k, v := range kv {
		keys = append(keys, k)
		values = append(values, v)
	}

	// TODO: Validate column name
	return NewDialect(t.driver).Insert(t.name, keys...).Values(values...)
}

// Update creates an update statement for the table name
func (t *Table) Update(kv map[string]interface{}) *Dialect {

	return NewDialect(t.driver).Update(t.Name()).Set(kv)
}

// Delete creates a delete statement for the table name
func (t *Table) Delete() *Dialect {

	return NewDialect(t.driver).Delete(t.Name())
}<commit_msg>reformat table<commit_after>package qb

import (
	"fmt"
	"strings"
)

// NewTable generates a new table pointer given table name, column and table constraints
func NewTable(driver string, name string, columns []Column, constraints []Constraint) *Table {
	return &Table{
		name:        name,
		columns:     columns,
		constraints: constraints,
		primaryCols: []string{},
		refs:        []ref{},
		driver:      driver,
	}
}

// Table is the base abstraction for any sql table
type Table struct {
	name        string
	columns     []Column
	constraints []Constraint
	primaryCols []string
	refs        []ref
	driver      string
}

// Name returns the table name
func (t *Table) Name() string {
	return t.name
}

// SQL generates create table syntax of table
func (t *Table) SQL() string {

	dialect := NewDialect(t.driver)

	cols := []string{}
	for _, v := range t.columns {
		cols = append(cols, v.SQL(t.driver))
	}

	constraints := []string{}

	// build primary key constraints using primaryCols
	if len(t.primaryCols) > 0 {
		constraints = append(constraints, fmt.Sprintf("PRIMARY KEY (%s)", strings.Join(t.primaryCols, ", ")))
	}

	// build foreign key constraints using refCols
	for _, ref := range t.refs {
		constraints = append(constraints, fmt.Sprintf("FOREIGN KEY (%s) REFERENCES %s(%s)", strings.Join(ref.cols, ", "), ref.refTable, strings.Join(ref.refCols, ", ")))
	}

	for _, v := range t.constraints {
		constraints = append(constraints, v.Name)
	}

	query := dialect.CreateTable(t.name, cols, constraints).Query()

	return query.SQL()
}

// AddColumn appends a new column to current table
func (t *Table) AddColumn(column Column) {
	t.columns = append(t.columns, column)
}

// AddConstraint appends a new constraint to current table
func (t *Table) AddConstraint(c Constraint) {
	t.constraints = append(t.constraints, c)
}

// AddPrimary appends a primary column that will be lazily built as a primary key constraint
func (t *Table) AddPrimary(col string) {
	t.primaryCols = append(t.primaryCols, col)
}

type ref struct {
	cols     []string
	refTable string
	refCols  []string
}

// AddRef appends a new reference struct that will be lazily built as a foreign key constraint
func (t *Table) AddRef(col string, refTable string, refCol string) {

	if len(t.refs) > 0 {
		for k, ref := range t.refs {
			if refTable == ref.refTable {
				t.refs[k].cols = append(t.refs[k].cols, fmt.Sprintf("%s", col))
				t.refs[k].refCols = append(t.refs[k].refCols, fmt.Sprintf("%s", refCol))
				return
			}
		}
	}

	r := ref{[]string{}, refTable, []string{}}
	r.cols = append(r.cols, col)
	r.refCols = append(r.refCols, refCol)
	t.refs = append(t.refs, r)
}

// Constraints returns the constraint slice of current table
func (t *Table) Constraints() []Constraint {
	return t.constraints
}

// Insert creates an insert statement for the table name
func (t *Table) Insert(kv map[string]interface{}) *Dialect {

	keys := []string{}
	values := []interface{}{}

	for k, v := range kv {
		keys = append(keys, k)
		values = append(values, v)
	}

	// TODO: Validate column name
	return NewDialect(t.driver).Insert(t.name, keys...).Values(values...)
}

// Update creates an update statement for the table name
func (t *Table) Update(kv map[string]interface{}) *Dialect {

	return NewDialect(t.driver).Update(t.Name()).Set(kv)
}

// Delete creates a delete statement for the table name
func (t *Table) Delete() *Dialect {
	return NewDialect(t.driver).Delete(t.Name())
}
<|endoftext|>


<commit_before>package exec

import (
	"bytes"
	"fmt"
	"os/exec"
	"path/filepath"
	"runtime"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/kballard/go-shellquote"

	"github.com/influxdata/telegraf"
	"github.com/influxdata/telegraf/internal"
	"github.com/influxdata/telegraf/plugins/inputs"
	"github.com/influxdata/telegraf/plugins/parsers"
	"github.com/influxdata/telegraf/plugins/parsers/nagios"
)

const sampleConfig = `
  ## Commands array
  commands = [
    "/tmp/test.sh",
    "/usr/bin/mycollector --foo=bar",
    "/tmp/collect_*.sh"
  ]

  ## Timeout for each command to complete.
  timeout = "5s"

  ## measurement name suffix (for separating different commands)
  name_suffix = "_mycollector"

  ## Data format to consume.
  ## Each data format has its own unique set of configuration options, read
  ## more about them here:
  ## https://github.com/influxdata/telegraf/blob/master/docs/DATA_FORMATS_INPUT.md
  data_format = "influx"
`

type Exec struct {
	Commands []string
	Command  string
	Timeout  internal.Duration

	parser parsers.Parser

	runner Runner
}

func NewExec() *Exec {
	return &Exec{
		runner:  CommandRunner{},
		Timeout: internal.Duration{Duration: time.Second * 5},
	}
}

type Runner interface {
	Run(*Exec, string, telegraf.Accumulator) ([]byte, error)
}

type CommandRunner struct{}

func AddNagiosState(exitCode error, acc telegraf.Accumulator) error {
	nagiosState := 0
	if exitCode != nil {
		exiterr, ok := exitCode.(*exec.ExitError)
		if ok {
			status, ok := exiterr.Sys().(syscall.WaitStatus)
			if ok {
				nagiosState = status.ExitStatus()
			} else {
				return fmt.Errorf("exec: unable to get nagios plugin exit code")
			}
		} else {
			return fmt.Errorf("exec: unable to get nagios plugin exit code")
		}
	}
	fields := map[string]interface{}{"state": nagiosState}
	acc.AddFields("nagios_state", fields, nil)
	return nil
}

func (c CommandRunner) Run(
	e *Exec,
	command string,
	acc telegraf.Accumulator,
) ([]byte, error) {
	split_cmd, err := shellquote.Split(command)
	if err != nil || len(split_cmd) == 0 {
		return nil, fmt.Errorf("exec: unable to parse command, %s", err)
	}

	cmd := exec.Command(split_cmd[0], split_cmd[1:]...)

	var out bytes.Buffer
	cmd.Stdout = &out

	if err := internal.RunTimeout(cmd, e.Timeout.Duration); err != nil {
		switch e.parser.(type) {
		case *nagios.NagiosParser:
			AddNagiosState(err, acc)
		default:
			return nil, fmt.Errorf("exec: %s for command '%s'", err, command)
		}
	} else {
		switch e.parser.(type) {
		case *nagios.NagiosParser:
			AddNagiosState(nil, acc)
		}
	}

	out = removeCarriageReturns(out)
	return out.Bytes(), nil
}

// removeCarriageReturns removes all carriage returns from the input if the
// OS is Windows. It does not return any errors.
func removeCarriageReturns(b bytes.Buffer) bytes.Buffer {
	if runtime.GOOS == "windows" {
		var buf bytes.Buffer
		for {
			byt, er := b.ReadBytes(0x0D)
			end := len(byt)
			if nil == er {
				end -= 1
			}
			if nil != byt {
				buf.Write(byt[:end])
			} else {
				break
			}
			if nil != er {
				break
			}
		}
		b = buf
	}
	return b

}

func (e *Exec) ProcessCommand(command string, acc telegraf.Accumulator, wg *sync.WaitGroup) {
	defer wg.Done()

	out, err := e.runner.Run(e, command, acc)
	if err != nil {
		acc.AddError(err)
		return
	}

	metrics, err := e.parser.Parse(out)
	if err != nil {
		acc.AddError(err)
	} else {
		for _, metric := range metrics {
			acc.AddFields(metric.Name(), metric.Fields(), metric.Tags(), metric.Time())
		}
	}
}

func (e *Exec) SampleConfig() string {
	return sampleConfig
}

func (e *Exec) Description() string {
	return "Read metrics from one or more commands that can output to stdout"
}

func (e *Exec) SetParser(parser parsers.Parser) {
	e.parser = parser
}

func (e *Exec) Gather(acc telegraf.Accumulator) error {
	var wg sync.WaitGroup
	// Legacy single command support
	if e.Command != "" {
		e.Commands = append(e.Commands, e.Command)
		e.Command = ""
	}

	commands := make([]string, 0, len(e.Commands))
	for _, pattern := range e.Commands {
		cmdAndArgs := strings.SplitN(pattern, " ", 2)
		if len(cmdAndArgs) == 0 {
			continue
		}

		matches, err := filepath.Glob(cmdAndArgs[0])
		if err != nil {
			acc.AddError(err)
			continue
		}

		if len(matches) == 0 {
			// There were no matches with the glob pattern, so let's assume
			// that the command is in PATH and just run it as it is
			commands = append(commands, pattern)
		} else {
			// There were matches, so we'll append each match together with
			// the arguments to the commands slice
			for _, match := range matches {
				if len(cmdAndArgs) == 1 {
					commands = append(commands, match)
				} else {
					commands = append(commands,
						strings.Join([]string{match, cmdAndArgs[1]}, " "))
				}
			}
		}
	}

	wg.Add(len(commands))
	for _, command := range commands {
		go e.ProcessCommand(command, acc, &wg)
	}
	wg.Wait()
	return nil
}

func init() {
	inputs.Add("exec", func() telegraf.Input {
		return NewExec()
	})
}
<commit_msg>Add output of stderr in case of error to exec log message (#3862)<commit_after>package exec

import (
	"bytes"
	"fmt"
	"os/exec"
	"path/filepath"
	"runtime"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/kballard/go-shellquote"

	"github.com/influxdata/telegraf"
	"github.com/influxdata/telegraf/internal"
	"github.com/influxdata/telegraf/plugins/inputs"
	"github.com/influxdata/telegraf/plugins/parsers"
	"github.com/influxdata/telegraf/plugins/parsers/nagios"
)

const sampleConfig = `
  ## Commands array
  commands = [
    "/tmp/test.sh",
    "/usr/bin/mycollector --foo=bar",
    "/tmp/collect_*.sh"
  ]

  ## Timeout for each command to complete.
  timeout = "5s"

  ## measurement name suffix (for separating different commands)
  name_suffix = "_mycollector"

  ## Data format to consume.
  ## Each data format has its own unique set of configuration options, read
  ## more about them here:
  ## https://github.com/influxdata/telegraf/blob/master/docs/DATA_FORMATS_INPUT.md
  data_format = "influx"
`

const MaxStderrBytes = 512

type Exec struct {
	Commands []string
	Command  string
	Timeout  internal.Duration

	parser parsers.Parser

	runner Runner
}

func NewExec() *Exec {
	return &Exec{
		runner:  CommandRunner{},
		Timeout: internal.Duration{Duration: time.Second * 5},
	}
}

type Runner interface {
	Run(*Exec, string, telegraf.Accumulator) ([]byte, error)
}

type CommandRunner struct{}

func AddNagiosState(exitCode error, acc telegraf.Accumulator) error {
	nagiosState := 0
	if exitCode != nil {
		exiterr, ok := exitCode.(*exec.ExitError)
		if ok {
			status, ok := exiterr.Sys().(syscall.WaitStatus)
			if ok {
				nagiosState = status.ExitStatus()
			} else {
				return fmt.Errorf("exec: unable to get nagios plugin exit code")
			}
		} else {
			return fmt.Errorf("exec: unable to get nagios plugin exit code")
		}
	}
	fields := map[string]interface{}{"state": nagiosState}
	acc.AddFields("nagios_state", fields, nil)
	return nil
}

func (c CommandRunner) Run(
	e *Exec,
	command string,
	acc telegraf.Accumulator,
) ([]byte, error) {
	split_cmd, err := shellquote.Split(command)
	if err != nil || len(split_cmd) == 0 {
		return nil, fmt.Errorf("exec: unable to parse command, %s", err)
	}

	cmd := exec.Command(split_cmd[0], split_cmd[1:]...)

	var (
		out    bytes.Buffer
		stderr bytes.Buffer
	)
	cmd.Stdout = &out
	cmd.Stderr = &stderr

	if err := internal.RunTimeout(cmd, e.Timeout.Duration); err != nil {
		switch e.parser.(type) {
		case *nagios.NagiosParser:
			AddNagiosState(err, acc)
		default:
			var errMessage = ""
			if stderr.Len() > 0 {
				stderr = removeCarriageReturns(stderr)
				// Limit the number of bytes.
				didTruncate := false
				if stderr.Len() > MaxStderrBytes {
					stderr.Truncate(MaxStderrBytes)
					didTruncate = true
				}
				if i := bytes.IndexByte(stderr.Bytes(), '\n'); i > 0 {
					// Only show truncation if the newline wasn't the last character.
					if i < stderr.Len()-1 {
						didTruncate = true
					}
					stderr.Truncate(i)
				}
				if didTruncate {
					stderr.WriteString("...")
				}

				errMessage = fmt.Sprintf(": %s", stderr.String())
			}
			return nil, fmt.Errorf("exec: %s for command '%s'%s", err, command, errMessage)
		}
	} else {
		switch e.parser.(type) {
		case *nagios.NagiosParser:
			AddNagiosState(nil, acc)
		}
	}

	out = removeCarriageReturns(out)
	return out.Bytes(), nil
}

// removeCarriageReturns removes all carriage returns from the input if the
// OS is Windows. It does not return any errors.
func removeCarriageReturns(b bytes.Buffer) bytes.Buffer {
	if runtime.GOOS == "windows" {
		var buf bytes.Buffer
		for {
			byt, er := b.ReadBytes(0x0D)
			end := len(byt)
			if nil == er {
				end -= 1
			}
			if nil != byt {
				buf.Write(byt[:end])
			} else {
				break
			}
			if nil != er {
				break
			}
		}
		b = buf
	}
	return b

}

func (e *Exec) ProcessCommand(command string, acc telegraf.Accumulator, wg *sync.WaitGroup) {
	defer wg.Done()

	out, err := e.runner.Run(e, command, acc)
	if err != nil {
		acc.AddError(err)
		return
	}

	metrics, err := e.parser.Parse(out)
	if err != nil {
		acc.AddError(err)
	} else {
		for _, metric := range metrics {
			acc.AddFields(metric.Name(), metric.Fields(), metric.Tags(), metric.Time())
		}
	}
}

func (e *Exec) SampleConfig() string {
	return sampleConfig
}

func (e *Exec) Description() string {
	return "Read metrics from one or more commands that can output to stdout"
}

func (e *Exec) SetParser(parser parsers.Parser) {
	e.parser = parser
}

func (e *Exec) Gather(acc telegraf.Accumulator) error {
	var wg sync.WaitGroup
	// Legacy single command support
	if e.Command != "" {
		e.Commands = append(e.Commands, e.Command)
		e.Command = ""
	}

	commands := make([]string, 0, len(e.Commands))
	for _, pattern := range e.Commands {
		cmdAndArgs := strings.SplitN(pattern, " ", 2)
		if len(cmdAndArgs) == 0 {
			continue
		}

		matches, err := filepath.Glob(cmdAndArgs[0])
		if err != nil {
			acc.AddError(err)
			continue
		}

		if len(matches) == 0 {
			// There were no matches with the glob pattern, so let's assume
			// that the command is in PATH and just run it as it is
			commands = append(commands, pattern)
		} else {
			// There were matches, so we'll append each match together with
			// the arguments to the commands slice
			for _, match := range matches {
				if len(cmdAndArgs) == 1 {
					commands = append(commands, match)
				} else {
					commands = append(commands,
						strings.Join([]string{match, cmdAndArgs[1]}, " "))
				}
			}
		}
	}

	wg.Add(len(commands))
	for _, command := range commands {
		go e.ProcessCommand(command, acc, &wg)
	}
	wg.Wait()
	return nil
}

func init() {
	inputs.Add("exec", func() telegraf.Input {
		return NewExec()
	})
}
<|endoftext|>


<commit_before>package nodejs_test

import (
	"os"

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

	"github.com/bouk/monkey"
	"github.com/markelog/archive"
	"github.com/markelog/cprf"
	"gopkg.in/cavaliercoder/grab.v1"

	"github.com/markelog/eclectica/io"
	. "github.com/markelog/eclectica/plugins/nodejs"
	"github.com/markelog/eclectica/variables"
)

var _ = Describe("yarn", func() {
	var (
		grabGet        bool
		archiveExtract bool
		cprfCopy       bool
		osRemovaAll    bool
		ioSymlink      bool
	)

	BeforeEach(func() {
		grabGet = false
		archiveExtract = false
		cprfCopy = false
		osRemovaAll = false
		ioSymlink = false

		monkey.Patch(grab.Get, func(path, url string) (*grab.Response, error) {
			grabGet = true
			return nil, nil
		})

		monkey.Patch(archive.Extract, func(from, to string) error {
			archiveExtract = true
			return nil
		})

		monkey.Patch(cprf.Copy, func(from, to string) error {
			cprfCopy = true
			return nil
		})

		monkey.Patch(os.RemoveAll, func(path string) error {
			osRemovaAll = true
			return nil
		})

		monkey.Patch(io.Symlink, func(from, to string) error {
			ioSymlink = true
			return nil
		})
	})

	AfterEach(func() {
		monkey.Unpatch(grab.Get)
		monkey.Unpatch(archive.Extract)
		monkey.Unpatch(cprf.Copy)
		monkey.Unpatch(os.RemoveAll)
		monkey.Unpatch(io.Symlink)
	})

	It("should not try to install yarn for unsupported node version", func() {
		working, err := (&Node{Version: "0.10.0"}).Yarn()

		Expect(grabGet).To(Equal(false))
		Expect(archiveExtract).To(Equal(false))
		Expect(cprfCopy).To(Equal(false))
		Expect(osRemovaAll).To(Equal(false))
		Expect(ioSymlink).To(Equal(false))

		Expect(working).To(Equal(true))
		Expect(err.Error()).To(Equal("\"0.10.0\" version is not supported by yarn"))
	})

	It("installs yarn", func() {
		monkey.Patch(variables.TempDir, func() string {
			path, _ := os.Getwd()

			return path
		})

		working, err := (&Node{Version: "6.10.0"}).Yarn()

		Expect(grabGet).To(Equal(true))
		Expect(archiveExtract).To(Equal(true))
		Expect(cprfCopy).To(Equal(true))
		Expect(osRemovaAll).To(Equal(true))
		Expect(ioSymlink).To(Equal(true))

		Expect(working).To(Equal(true))
		Expect(err).To(BeNil())

		monkey.Unpatch(variables.TempDir())
	})
})
<commit_msg>Remove test for removeAll<commit_after>package nodejs_test

import (
	"os"

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

	"github.com/bouk/monkey"
	"github.com/markelog/archive"
	"github.com/markelog/cprf"
	"gopkg.in/cavaliercoder/grab.v1"

	"github.com/markelog/eclectica/io"
	. "github.com/markelog/eclectica/plugins/nodejs"
	"github.com/markelog/eclectica/variables"
)

var _ = Describe("yarn", func() {
	var (
		grabGet        bool
		archiveExtract bool
		cprfCopy       bool
		ioSymlink      bool
	)

	BeforeEach(func() {
		grabGet = false
		archiveExtract = false
		cprfCopy = false
		ioSymlink = false

		monkey.Patch(grab.Get, func(path, url string) (*grab.Response, error) {
			grabGet = true
			return nil, nil
		})

		monkey.Patch(archive.Extract, func(from, to string) error {
			archiveExtract = true
			return nil
		})

		monkey.Patch(cprf.Copy, func(from, to string) error {
			cprfCopy = true
			return nil
		})

		monkey.Patch(io.Symlink, func(from, to string) error {
			ioSymlink = true
			return nil
		})
	})

	AfterEach(func() {
		monkey.Unpatch(grab.Get)
		monkey.Unpatch(archive.Extract)
		monkey.Unpatch(cprf.Copy)
		monkey.Unpatch(io.Symlink)
	})

	It("should not try to install yarn for unsupported node version", func() {
		working, err := (&Node{Version: "0.10.0"}).Yarn()

		Expect(grabGet).To(Equal(false))
		Expect(archiveExtract).To(Equal(false))
		Expect(cprfCopy).To(Equal(false))
		Expect(ioSymlink).To(Equal(false))

		Expect(working).To(Equal(true))
		Expect(err.Error()).To(Equal("\"0.10.0\" version is not supported by yarn"))
	})

	It("installs yarn", func() {
		monkey.Patch(variables.TempDir, func() string {
			path, _ := os.Getwd()

			return path
		})

		working, err := (&Node{Version: "6.10.0"}).Yarn()

		Expect(grabGet).To(Equal(true))
		Expect(archiveExtract).To(Equal(true))
		Expect(cprfCopy).To(Equal(true))
		Expect(ioSymlink).To(Equal(true))

		Expect(working).To(Equal(true))
		Expect(err).To(BeNil())

		monkey.Unpatch(variables.TempDir())
	})
})
<|endoftext|>


<commit_before>// Package untappd implements a plugin to get info from Untappd users.
package untappd

import (
	"errors"
	"fmt"
	"io/ioutil"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/StalkR/goircbot/bot"
	"github.com/StalkR/goircbot/lib/duration"
	"github.com/StalkR/goircbot/lib/transport"
)

var (
	indexRE   = regexp.MustCompile(`<title>Untappd`)
	nameRE    = regexp.MustCompile(`<title>(.*) on Untappd</title>`)
	totalRE   = regexp.MustCompile(`(?s)<p><strong>Total</strong></p>\s*<span>(\d+)`)
	detailsRE = regexp.MustCompile(`(?s)<div class="details">(.*)`)
	beerRE    = regexp.MustCompile(`(?s)<a href=./b/[^>]+>([^<]+)`)
	whenRE    = regexp.MustCompile(`(?s)<li class="time[^>]+>([^<]+)</li>`)
)

// Info represents an Untappd user info with last activity and total.
type Info struct {
	Name  string
	Beer  string
	When  time.Time
	Total int
}

// String formats an Info to fit on a single IRC line.
func (i Info) String() string {
	if i.Total == 0 {
		return fmt.Sprintf("%v doesn't drink, booo!", i.Name)
	}
	return fmt.Sprintf("%v is drinking %v (%v ago) - Total beers: %v",
		i.Name, i.Beer, duration.Format(time.Since(i.When)), i.Total)
}

// userPage fetches an Untappd user page.
func userPage(user string) ([]byte, error) {
	url := fmt.Sprintf("https://untappd.com/user/%s", user)
	client, err := transport.Client(url)
	if err != nil {
		return nil, err
	}
	resp, err := client.Get(url)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	page, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}
	return page, nil
}

// parsePage parses an Untappd user page and return its info.
func parsePage(page []byte) (i Info, err error) {
	if indexRE.Match(page) {
		return i, errors.New("no such user")
	}

	match := nameRE.FindSubmatch(page)
	if match == nil {
		return i, errors.New("untappd: could not parse name")
	}
	i.Name = strings.TrimSpace(string(match[1]))

	match = totalRE.FindSubmatch(page)
	if match == nil {
		return i, errors.New("untappd: could not parse total")
	}
	i.Total, err = strconv.Atoi(strings.TrimSpace(string(match[1])))
	if err != nil {
		return i, fmt.Errorf("untappd: error parsing total: %v", err)
	}
	if i.Total == 0 {
		return i, nil
	}

	match = detailsRE.FindSubmatch(page)
	if match == nil {
		return i, errors.New("untappd: could not parse details")
	}
	details := match[1]

	match = beerRE.FindSubmatch(details)
	if match == nil {
		return i, errors.New("untappd: could not parse beer")
	}
	i.Beer = strings.TrimSpace(string(match[1]))

	match = whenRE.FindSubmatch(details)
	if match == nil {
		return i, errors.New("untappd: could not parse when")
	}
	when, err := time.Parse(time.RFC1123Z, strings.TrimSpace(string(match[1])))
	if err != nil {
		return i, fmt.Errorf("untappd: error parsing time: %v", err)
	}
	i.When = when

	return
}

// UserInfo obtains info for a user by querying its web page and parsing it.
func UserInfo(user string) (Info, error) {
	page, err := userPage(user)
	if err != nil {
		return Info{}, fmt.Errorf("untappd: error getting page: %v", err)
	}
	return parsePage(page)
}

// untappd handles an IRC command to print an Untappd user info.
func untappd(b *bot.Bot, e *bot.Event) {
	user := strings.TrimSpace(e.Args)
	if len(user) == 0 {
		return
	}
	info, err := UserInfo(user)
	if err != nil {
		b.Conn.Privmsg(e.Target, err.Error())
		return
	}
	b.Conn.Privmsg(e.Target, info.String())
}

// Register registers the plugin with a bot.
func Register(b *bot.Bot) {
	b.AddCommand("untappd", bot.Command{
		Help:    "get stats and activity from an Untappd user",
		Handler: untappd,
		Pub:     true,
		Priv:    true,
		Hidden:  false})
}
<commit_msg>plugins/untappd: rewrite with new bot<commit_after>// Package untappd implements a plugin to get info from Untappd users.
package untappd

import (
	"errors"
	"fmt"
	"io/ioutil"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/StalkR/goircbot/bot"
	"github.com/StalkR/goircbot/lib/duration"
	"github.com/StalkR/goircbot/lib/transport"
)

var (
	indexRE   = regexp.MustCompile(`<title>Untappd`)
	nameRE    = regexp.MustCompile(`<title>(.*) on Untappd</title>`)
	totalRE   = regexp.MustCompile(`(?s)<p><strong>Total</strong></p>\s*<span>(\d+)`)
	detailsRE = regexp.MustCompile(`(?s)<div class="details">(.*)`)
	beerRE    = regexp.MustCompile(`(?s)<a href=./b/[^>]+>([^<]+)`)
	whenRE    = regexp.MustCompile(`(?s)<li class="time[^>]+>([^<]+)</li>`)
)

// Info represents an Untappd user info with last activity and total.
type Info struct {
	Name  string
	Beer  string
	When  time.Time
	Total int
}

// String formats an Info to fit on a single IRC line.
func (i Info) String() string {
	if i.Total == 0 {
		return fmt.Sprintf("%v doesn't drink, booo!", i.Name)
	}
	return fmt.Sprintf("%v is drinking %v (%v ago) - Total beers: %v",
		i.Name, i.Beer, duration.Format(time.Since(i.When)), i.Total)
}

// userPage fetches an Untappd user page.
func userPage(user string) ([]byte, error) {
	url := fmt.Sprintf("https://untappd.com/user/%s", user)
	client, err := transport.Client(url)
	if err != nil {
		return nil, err
	}
	resp, err := client.Get(url)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	page, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}
	return page, nil
}

// parsePage parses an Untappd user page and return its info.
func parsePage(page []byte) (i Info, err error) {
	if indexRE.Match(page) {
		return i, errors.New("no such user")
	}

	match := nameRE.FindSubmatch(page)
	if match == nil {
		return i, errors.New("untappd: could not parse name")
	}
	i.Name = strings.TrimSpace(string(match[1]))

	match = totalRE.FindSubmatch(page)
	if match == nil {
		return i, errors.New("untappd: could not parse total")
	}
	i.Total, err = strconv.Atoi(strings.TrimSpace(string(match[1])))
	if err != nil {
		return i, fmt.Errorf("untappd: error parsing total: %v", err)
	}
	if i.Total == 0 {
		return i, nil
	}

	match = detailsRE.FindSubmatch(page)
	if match == nil {
		return i, errors.New("untappd: could not parse details")
	}
	details := match[1]

	match = beerRE.FindSubmatch(details)
	if match == nil {
		return i, errors.New("untappd: could not parse beer")
	}
	i.Beer = strings.TrimSpace(string(match[1]))

	match = whenRE.FindSubmatch(details)
	if match == nil {
		return i, errors.New("untappd: could not parse when")
	}
	when, err := time.Parse(time.RFC1123Z, strings.TrimSpace(string(match[1])))
	if err != nil {
		return i, fmt.Errorf("untappd: error parsing time: %v", err)
	}
	i.When = when

	return
}

// UserInfo obtains info for a user by querying its web page and parsing it.
func UserInfo(user string) (Info, error) {
	page, err := userPage(user)
	if err != nil {
		return Info{}, fmt.Errorf("untappd: error getting page: %v", err)
	}
	return parsePage(page)
}

// untappd handles an IRC command to print an Untappd user info.
func untappd(e *bot.Event) {
	user := strings.TrimSpace(e.Args)
	if len(user) == 0 {
		return
	}
	info, err := UserInfo(user)
	if err != nil {
		e.Bot.Privmsg(e.Target, err.Error())
		return
	}
	e.Bot.Privmsg(e.Target, info.String())
}

// Register registers the plugin with a bot.
func Register(b bot.Bot) {
	b.Commands().Add("untappd", bot.Command{
		Help:    "get stats and activity from an Untappd user",
		Handler: untappd,
		Pub:     true,
		Priv:    true,
		Hidden:  false})
}
<|endoftext|>


<commit_before>package box

import (
	"bytes"
	"encoding/json"
	"fmt"
	"net/http"
)

type TaskCollection struct {
	TotalCount int     `json:"total_count"`
	Entries    []*Task `json:"entries"`
}

// Documentation: https://developers.box.com/docs/#tasks-task-object
// TODO(ttacon): add missing fields
type Task struct {
	Type                     string                    `json:"type"`
	Id                       string                    `json:"id"`
	Item                     *Item                     `json:"item"`
	DueAt                    *string                   `json:"due_at"`     // TODO(ttacon): time.Time
	CreatedAt                *string                   `json:"created_at"` // TODO(ttacon): time.Time
	CreatedBy                *Item                     `json:"created_by"` // TODO(ttacon): change to user
	Action                   *string                   `json:"action"`     //TODO(ttacon): validation as this must be 'review'?
	Message                  *string                   `json:"message"`
	IsCompleted              *bool                     `json:"is_completed"`
	TaskAssignmentCollection *TaskAssignmentCollection `json:"task_assignment_collection"`
}

type TaskAssignmentCollection struct {
	TotalCount int               `json:"total_count"`
	Entries    []*TaskAssignment `json:"entries"`
}

// TODO(ttacon): find out where the deuce this is defined in their documentation?!?!?!
type TaskAssignment struct {
	Type       *string `json:"type"`
	Id         string  `json:"id"`
	Item       *Item   `json:"item"`
	AssignedTo *Item   `json:"assigned_to"` // TODO(ttacon): change to mini-user
}

// Documentation: https://developers.box.com/docs/#tasks-create-a-task
func (c *Client) CreateTask(itemId, itemType, action, message, due_at string) (*http.Response, *Task, error) {
	var dataMap = map[string]interface{}{
		"item": map[string]string{
			"id":   itemId,
			"type": itemType,
		},
	}
	if len(action) > 0 {
		// TODO(ttacon): make sure this is "review"
		dataMap["action"] = action
	}
	if len(message) > 0 {
		dataMap["message"] = message
	}
	if len(due_at) > 0 {
		dataMap["due_at"] = due_at
	}

	dataBytes, err := json.Marshal(dataMap)
	if err != nil {
		return nil, nil, err
	}

	req, err := http.NewRequest(
		"POST",
		fmt.Sprintf("%s/tasks", BASE_URL),
		bytes.NewReader(dataBytes),
	)
	if err != nil {
		return nil, nil, err
	}

	resp, err := c.Trans.Client().Do(req)
	if err != nil {
		return resp, nil, err
	}

	var data Task
	err = json.NewDecoder(resp.Body).Decode(&data)
	return resp, &data, err
}

// Documentation: https://developers.box.com/docs/#tasks-get-a-task
func (c *Client) GetTask(taskId string) (*http.Response, *Task, error) {
	req, err := http.NewRequest(
		"GET",
		fmt.Sprintf("%s/tasks/%s", BASE_URL, taskId),
		nil,
	)
	if err != nil {
		return nil, nil, err
	}

	resp, err := c.Trans.Client().Do(req)
	if err != nil {
		return resp, nil, err
	}

	var data Task
	err = json.NewDecoder(resp.Body).Decode(&data)
	return resp, &data, err
}

// Documentation: https://developers.box.com/docs/#tasks-update-a-task
func (c *Client) UpdateTask(taskId, action, message, due_at string) (*http.Response, *Task, error) {
	var dataMap = make(map[string]interface{})
	if len(action) > 0 {
		dataMap["action"] = action
	}
	if len(message) > 0 {
		dataMap["message"] = message
	}
	if len(due_at) > 0 {
		dataMap["due_at"] = due_at
	}

	dataBytes, err := json.Marshal(dataMap)
	if err != nil {
		return nil, nil, err
	}

	req, err := http.NewRequest(
		"PUT",
		fmt.Sprintf("%s/tasks/%s", BASE_URL, taskId),
		bytes.NewReader(dataBytes),
	)
	if err != nil {
		return nil, nil, err
	}

	resp, err := c.Trans.Client().Do(req)
	if err != nil {
		return resp, nil, err
	}

	var data Task
	err = json.NewDecoder(resp.Body).Decode(&data)
	return resp, &data, err
}

// Documentation: https://developers.box.com/docs/#tasks-delete-a-task
func (c *Client) DeleteTask(taskId string) (*http.Response, error) {
	req, err := http.NewRequest(
		"DELETE",
		fmt.Sprintf("%s/tasks/%s", BASE_URL, taskId),
		nil,
	)
	if err != nil {
		return nil, nil, err
	}

	return c.Trans.Client().Do(req)
}
<commit_msg>Add get task assignments<commit_after>package box

import (
	"bytes"
	"encoding/json"
	"fmt"
	"net/http"
)

type TaskCollection struct {
	TotalCount int     `json:"total_count"`
	Entries    []*Task `json:"entries"`
}

// Documentation: https://developers.box.com/docs/#tasks-task-object
// TODO(ttacon): add missing fields
type Task struct {
	Type                     string                    `json:"type"`
	Id                       string                    `json:"id"`
	Item                     *Item                     `json:"item"`
	DueAt                    *string                   `json:"due_at"`     // TODO(ttacon): time.Time
	CreatedAt                *string                   `json:"created_at"` // TODO(ttacon): time.Time
	CreatedBy                *Item                     `json:"created_by"` // TODO(ttacon): change to user
	Action                   *string                   `json:"action"`     //TODO(ttacon): validation as this must be 'review'?
	Message                  *string                   `json:"message"`
	IsCompleted              *bool                     `json:"is_completed"`
	TaskAssignmentCollection *TaskAssignmentCollection `json:"task_assignment_collection"`
}

type TaskAssignmentCollection struct {
	TotalCount int               `json:"total_count"`
	Entries    []*TaskAssignment `json:"entries"`
}

// TODO(ttacon): find out where the deuce this is defined in their documentation?!?!?!
type TaskAssignment struct {
	Type       *string `json:"type"`
	Id         string  `json:"id"`
	Item       *Item   `json:"item"`
	AssignedTo *Item   `json:"assigned_to"` // TODO(ttacon): change to mini-user
}

// Documentation: https://developers.box.com/docs/#tasks-create-a-task
func (c *Client) CreateTask(itemId, itemType, action, message, due_at string) (*http.Response, *Task, error) {
	var dataMap = map[string]interface{}{
		"item": map[string]string{
			"id":   itemId,
			"type": itemType,
		},
	}
	if len(action) > 0 {
		// TODO(ttacon): make sure this is "review"
		dataMap["action"] = action
	}
	if len(message) > 0 {
		dataMap["message"] = message
	}
	if len(due_at) > 0 {
		dataMap["due_at"] = due_at
	}

	dataBytes, err := json.Marshal(dataMap)
	if err != nil {
		return nil, nil, err
	}

	req, err := http.NewRequest(
		"POST",
		fmt.Sprintf("%s/tasks", BASE_URL),
		bytes.NewReader(dataBytes),
	)
	if err != nil {
		return nil, nil, err
	}

	resp, err := c.Trans.Client().Do(req)
	if err != nil {
		return resp, nil, err
	}

	var data Task
	err = json.NewDecoder(resp.Body).Decode(&data)
	return resp, &data, err
}

// Documentation: https://developers.box.com/docs/#tasks-get-a-task
func (c *Client) GetTask(taskId string) (*http.Response, *Task, error) {
	req, err := http.NewRequest(
		"GET",
		fmt.Sprintf("%s/tasks/%s", BASE_URL, taskId),
		nil,
	)
	if err != nil {
		return nil, nil, err
	}

	resp, err := c.Trans.Client().Do(req)
	if err != nil {
		return resp, nil, err
	}

	var data Task
	err = json.NewDecoder(resp.Body).Decode(&data)
	return resp, &data, err
}

// Documentation: https://developers.box.com/docs/#tasks-update-a-task
func (c *Client) UpdateTask(taskId, action, message, due_at string) (*http.Response, *Task, error) {
	var dataMap = make(map[string]interface{})
	if len(action) > 0 {
		dataMap["action"] = action
	}
	if len(message) > 0 {
		dataMap["message"] = message
	}
	if len(due_at) > 0 {
		dataMap["due_at"] = due_at
	}

	dataBytes, err := json.Marshal(dataMap)
	if err != nil {
		return nil, nil, err
	}

	req, err := http.NewRequest(
		"PUT",
		fmt.Sprintf("%s/tasks/%s", BASE_URL, taskId),
		bytes.NewReader(dataBytes),
	)
	if err != nil {
		return nil, nil, err
	}

	resp, err := c.Trans.Client().Do(req)
	if err != nil {
		return resp, nil, err
	}

	var data Task
	err = json.NewDecoder(resp.Body).Decode(&data)
	return resp, &data, err
}

// Documentation: https://developers.box.com/docs/#tasks-delete-a-task
func (c *Client) DeleteTask(taskId string) (*http.Response, error) {
	req, err := http.NewRequest(
		"DELETE",
		fmt.Sprintf("%s/tasks/%s", BASE_URL, taskId),
		nil,
	)
	if err != nil {
		return nil, nil, err
	}

	return c.Trans.Client().Do(req)
}

// Documentation: https://developers.box.com/docs/#tasks-get-the-assignments-for-a-task
func (c *Client) GetAssignmentsForTask(taskId string) (*http.Response, *TaskAssignmentCollection, error) {
	req, err := http.NewRequest(
		"GET",
		fmt.Sprintf("%s/tasks/%s/assignments", BASE_URL, taskId),
		nil,
	)
	if err != nil {
		return nil, nil, err
	}

	resp, err := c.Trans.Client().Do(req)
	if err != nil {
		return resp, nil, err
	}

	var data TaskAssignmentCollection
	err = json.NewDecoder(resp.Body).Decode(&data)
	return resp, &data, err
}
<|endoftext|>


<commit_before>package importer

import (
	"reflect"
	"strings"

	"github.com/tcard/sgo/sgo/annotations"
	"github.com/tcard/sgo/sgo/ast"
	"github.com/tcard/sgo/sgo/parser"
	"github.com/tcard/sgo/sgo/types"
)

// ConvertAST transforms an imported Go AST applying type annotations found in
// doc comments.
//
// It looks up doc comment lines of the form 'For SGO: <annotation>',
// where <annotation> is a type expression. If found, the annotation replaces
// the type of the documented declaration or field.
func ConvertAST(a *ast.File, info *types.Info, ann *annotations.Annotation) {
	c := astConverter{info: info}
	c.convertAST(a, ann, nil)
}

type astConverter struct {
	info      *types.Info
	converted map[interface{}]struct{}
}

func (c *astConverter) convertAST(node ast.Node, ann *annotations.Annotation, replace func(e ast.Expr)) {
	if replaced := c.maybeReplace(node, ann, replace); replaced {
		return
	}

	switch n := node.(type) {
	case *ast.Field:
		if len(n.Names) == 1 {
			ann = ann.Lookup(n.Names[0].Name)
		} else {
			ann = nil
		}
		c.convertAST(n.Type, ann, func(e ast.Expr) { n.Type = e })

	case *ast.FieldList:
		for _, f := range n.List {
			c.convertAST(f, ann, nil)
		}

	case *ast.StarExpr:
		if replace != nil {
			replace(&ast.OptionalType{Elt: n})
		}
		c.convertAST(n.X, ann, func(e ast.Expr) { n.X = e })

	case *ast.Ident:
		u, ok := c.info.Uses[n]
		if !ok {
			break
		}
		tn, ok := u.(*types.TypeName)
		if !ok {
			break
		}
		if replace != nil && types.IsInterface(tn.Type()) {
			replace(&ast.OptionalType{Elt: n})
		}

	// Types
	case *ast.ArrayType:
		c.convertAST(n.Elt, ann, func(e ast.Expr) { n.Elt = e })

	case *ast.StructType:
		c.convertAST(n.Fields, ann, nil)

	case *ast.FuncType:
		if replace != nil {
			replace(&ast.OptionalType{Elt: n})
		}
		if n.Params != nil {
			c.convertAST(n.Params, ann, nil)
		}
		if n.Results != nil {
			c.convertAST(n.Results, ann, nil)
		}

	case *ast.InterfaceType:
		if replace != nil {
			replace(&ast.OptionalType{Elt: n})
		}
		for _, f := range n.Methods.List {
			name := ""
			if len(f.Names) == 0 {
				var id *ast.Ident

				switch t := f.Type.(type) {
				case *ast.Ident:
					id = t
				case *ast.StarExpr:
					if t, ok := t.X.(*ast.Ident); ok {
						id = t
					}
				}

				if id != nil {
					name = id.Name
				}
			} else {
				name = f.Names[0].Name
			}
			// Call maybeReplace here because, if it won't replace anything, we don't
			// want to pass a replace function to convertAST as it would think that
			// we're converting a function literal and make it optional by default.
			if replaced := c.maybeReplace(f.Type, ann.Lookup(name), func(e ast.Expr) { f.Type = e }); replaced {
				return
			}
			c.convertAST(f.Type, ann.Lookup(name), nil)
		}

	case *ast.MapType:
		if replace != nil {
			replace(&ast.OptionalType{Elt: n})
		}
		c.convertAST(n.Key, ann, func(e ast.Expr) { n.Key = e })
		c.convertAST(n.Value, ann, func(e ast.Expr) { n.Value = e })

	case *ast.ChanType:
		if replace != nil {
			replace(&ast.OptionalType{Elt: n})
		}
		c.convertAST(n.Value, ann, func(e ast.Expr) { n.Value = e })

	// Declarations
	case *ast.ValueSpec:
		if n.Type != nil {
			c.convertAST(n.Type, nil, func(e ast.Expr) { n.Type = e })
		}

	case *ast.TypeSpec:
		if n.Type != nil {
			// c.convertAST(n.Type, ann, func(e ast.Expr) { n.Type = e })
			c.convertAST(n.Type, ann, nil)
		}

	case *ast.GenDecl:
		for _, s := range n.Specs {
			switch s := s.(type) {
			case *ast.ImportSpec:
				c.convertAST(s, nil, nil)
			case *ast.ValueSpec:
				if replace == nil { // First time.
					c.convertAST(n, ann.Lookup(s.Names.List[0].Name), func(e ast.Expr) { s.Type = e })
				} else { // Second time.
					c.convertAST(s, ann, replace)
				}
			case *ast.TypeSpec:
				if replace == nil { // First time.
					c.convertAST(n, ann.Lookup(s.Name.Name), func(e ast.Expr) { s.Type = e })
				} else { // Second time.
					c.convertAST(s, ann, replace)
				}
			}
		}

	case *ast.FuncDecl:
		// https://github.com/tcard/sgo/issues/13
		// if n.Recv != nil {
		// 	c.convertAST(n.Recv, nil)
		// }

		// Call maybeReplace here because, if it won't replace anything, we don't
		// want to pass a replace function to convertAST as it would think that
		// we're converting a function literal and make it optional by default.
		if replaced := c.maybeReplace(n.Type, ann, func(e ast.Expr) { n.Type = e.(*ast.FuncType) }); replaced {
			return
		}
		c.convertAST(n.Type, nil, nil)

	case *ast.File:
		for _, d := range n.Decls {
			switch d := d.(type) {
			case *ast.GenDecl:
				c.convertAST(d, ann, nil)
			case *ast.FuncDecl:
				name := d.Name.Name
				if d.Recv != nil && len(d.Recv.List) > 0 {
					switch t := d.Recv.List[0].Type.(type) {
					case *ast.StarExpr:
						if id, ok := t.X.(*ast.Ident); ok {
							name = "(*" + id.Name + ")." + name
						}
					case *ast.Ident:
						name = t.Name + "." + name
					}
				}
				c.convertAST(d, ann.Lookup(name), func(e ast.Expr) {
					if e, ok := e.(*ast.FuncType); ok {
						d.Type = e
					}
				})
			}
		}
	}
}

func (c *astConverter) maybeReplace(node ast.Node, ann *annotations.Annotation, replace func(e ast.Expr)) bool {
	if replace == nil {
		return false
	}

	if typ, ok := ann.Type(); ok {
		e, err := parser.ParseExpr(typ)
		if err == nil {
			replace(e)
			return true
		}
	}

	n := reflect.ValueOf(node)

	if n.Elem().Type().Kind() != reflect.Struct {
		return false
	}

	doc := n.Elem().FieldByName("Doc")
	if !doc.IsValid() {
		return false
	}

	cg := doc.Interface().(*ast.CommentGroup)
	if cg == nil {
		return false
	}

	for _, l := range cg.List {
		s := l.Text
		s = strings.TrimPrefix(s, "//")
		s = strings.TrimPrefix(s, "/*")
		s = strings.TrimSpace(s)

		if !strings.HasPrefix(s, "For SGo: ") {
			continue
		}

		ann := s[len("For SGo: "):]
		e, err := parser.ParseExpr(ann)
		if err != nil {
			return false
		}

		replace(e)
		return true
	}

	return false
}
<commit_msg>Convert optionable method receivers to optional.<commit_after>package importer

import (
	"reflect"
	"strings"

	"github.com/tcard/sgo/sgo/annotations"
	"github.com/tcard/sgo/sgo/ast"
	"github.com/tcard/sgo/sgo/parser"
	"github.com/tcard/sgo/sgo/types"
)

// ConvertAST transforms an imported Go AST applying type annotations found in
// doc comments.
//
// It looks up doc comment lines of the form 'For SGO: <annotation>',
// where <annotation> is a type expression. If found, the annotation replaces
// the type of the documented declaration or field.
func ConvertAST(a *ast.File, info *types.Info, ann *annotations.Annotation) {
	c := astConverter{info: info}
	c.convertAST(a, ann, nil)
}

type astConverter struct {
	info      *types.Info
	converted map[interface{}]struct{}
}

func (c *astConverter) convertAST(node ast.Node, ann *annotations.Annotation, replace func(e ast.Expr)) {
	if replaced := c.maybeReplace(node, ann, replace); replaced {
		return
	}

	switch n := node.(type) {
	case *ast.Field:
		if len(n.Names) == 1 {
			ann = ann.Lookup(n.Names[0].Name)
		} else {
			ann = nil
		}
		c.convertAST(n.Type, ann, func(e ast.Expr) { n.Type = e })

	case *ast.FieldList:
		for _, f := range n.List {
			c.convertAST(f, ann, nil)
		}

	case *ast.StarExpr:
		if replace != nil {
			replace(&ast.OptionalType{Elt: n})
		}
		c.convertAST(n.X, ann, func(e ast.Expr) { n.X = e })

	case *ast.Ident:
		u, ok := c.info.Uses[n]
		if !ok {
			break
		}
		tn, ok := u.(*types.TypeName)
		if !ok {
			break
		}
		if replace != nil && types.IsOptionable(tn.Type()) {
			replace(&ast.OptionalType{Elt: n})
		}

	// Types
	case *ast.ArrayType:
		c.convertAST(n.Elt, ann, func(e ast.Expr) { n.Elt = e })

	case *ast.StructType:
		c.convertAST(n.Fields, ann, nil)

	case *ast.FuncType:
		if replace != nil {
			replace(&ast.OptionalType{Elt: n})
		}
		if n.Params != nil {
			c.convertAST(n.Params, ann, nil)
		}
		if n.Results != nil {
			c.convertAST(n.Results, ann, nil)
		}

	case *ast.InterfaceType:
		if replace != nil {
			replace(&ast.OptionalType{Elt: n})
		}
		for _, f := range n.Methods.List {
			name := ""
			if len(f.Names) == 0 {
				var id *ast.Ident

				switch t := f.Type.(type) {
				case *ast.Ident:
					id = t
				case *ast.StarExpr:
					if t, ok := t.X.(*ast.Ident); ok {
						id = t
					}
				}

				if id != nil {
					name = id.Name
				}
			} else {
				name = f.Names[0].Name
			}
			// Call maybeReplace here because, if it won't replace anything, we don't
			// want to pass a replace function to convertAST as it would think that
			// we're converting a function literal and make it optional by default.
			if replaced := c.maybeReplace(f.Type, ann.Lookup(name), func(e ast.Expr) { f.Type = e }); replaced {
				return
			}
			c.convertAST(f.Type, ann.Lookup(name), nil)
		}

	case *ast.MapType:
		if replace != nil {
			replace(&ast.OptionalType{Elt: n})
		}
		c.convertAST(n.Key, ann, func(e ast.Expr) { n.Key = e })
		c.convertAST(n.Value, ann, func(e ast.Expr) { n.Value = e })

	case *ast.ChanType:
		if replace != nil {
			replace(&ast.OptionalType{Elt: n})
		}
		c.convertAST(n.Value, ann, func(e ast.Expr) { n.Value = e })

	// Declarations
	case *ast.ValueSpec:
		if n.Type != nil {
			c.convertAST(n.Type, nil, func(e ast.Expr) { n.Type = e })
		}

	case *ast.TypeSpec:
		if n.Type != nil {
			// c.convertAST(n.Type, ann, func(e ast.Expr) { n.Type = e })
			c.convertAST(n.Type, ann, nil)
		}

	case *ast.GenDecl:
		for _, s := range n.Specs {
			switch s := s.(type) {
			case *ast.ImportSpec:
				c.convertAST(s, nil, nil)
			case *ast.ValueSpec:
				if replace == nil { // First time.
					c.convertAST(n, ann.Lookup(s.Names.List[0].Name), func(e ast.Expr) { s.Type = e })
				} else { // Second time.
					c.convertAST(s, ann, replace)
				}
			case *ast.TypeSpec:
				if replace == nil { // First time.
					c.convertAST(n, ann.Lookup(s.Name.Name), func(e ast.Expr) { s.Type = e })
				} else { // Second time.
					c.convertAST(s, ann, replace)
				}
			}
		}

	case *ast.FuncDecl:
		if n.Recv != nil {
			recv := n.Recv.List[0]
			switch typ := recv.Type.(type) {
			case *ast.StarExpr:
				recv.Type = &ast.OptionalType{Elt: typ}
			case *ast.Ident:
				c.convertAST(recv.Type, nil, func(e ast.Expr) { recv.Type = e })
			}
		}

		// Call maybeReplace here because, if it won't replace anything, we don't
		// want to pass a replace function to convertAST as it would think that
		// we're converting a function literal and make it optional by default.
		if replaced := c.maybeReplace(n.Type, ann, func(e ast.Expr) { n.Type = e.(*ast.FuncType) }); replaced {
			return
		}
		c.convertAST(n.Type, nil, nil)

	case *ast.File:
		for _, d := range n.Decls {
			switch d := d.(type) {
			case *ast.GenDecl:
				c.convertAST(d, ann, nil)
			case *ast.FuncDecl:
				name := d.Name.Name
				if d.Recv != nil && len(d.Recv.List) > 0 {
					switch t := d.Recv.List[0].Type.(type) {
					case *ast.StarExpr:
						if id, ok := t.X.(*ast.Ident); ok {
							name = "(*" + id.Name + ")." + name
						}
					case *ast.Ident:
						name = t.Name + "." + name
					}
				}
				c.convertAST(d, ann.Lookup(name), func(e ast.Expr) {
					if e, ok := e.(*ast.FuncType); ok {
						d.Type = e
					}
				})
			}
		}
	}
}

func (c *astConverter) maybeReplace(node ast.Node, ann *annotations.Annotation, replace func(e ast.Expr)) bool {
	if replace == nil {
		return false
	}

	if typ, ok := ann.Type(); ok {
		e, err := parser.ParseExpr(typ)
		if err == nil {
			replace(e)
			return true
		}
	}

	n := reflect.ValueOf(node)

	if n.Elem().Type().Kind() != reflect.Struct {
		return false
	}

	doc := n.Elem().FieldByName("Doc")
	if !doc.IsValid() {
		return false
	}

	cg := doc.Interface().(*ast.CommentGroup)
	if cg == nil {
		return false
	}

	for _, l := range cg.List {
		s := l.Text
		s = strings.TrimPrefix(s, "//")
		s = strings.TrimPrefix(s, "/*")
		s = strings.TrimSpace(s)

		if !strings.HasPrefix(s, "For SGo: ") {
			continue
		}

		ann := s[len("For SGo: "):]
		e, err := parser.ParseExpr(ann)
		if err != nil {
			return false
		}

		replace(e)
		return true
	}

	return false
}
<|endoftext|>


<commit_before>package main

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

var colors = []struct {
	desc string
	rgb  uint32
	hsl  uint32
}{
	{
		desc: "black",
		rgb:  0xFF000000,
		hsl:  0xFF000000,
	},
	{
		desc: "red",
		rgb:  0xFFFF0000,
		hsl:  0xFF00FF7F,
	},
	{
		desc: "yellow",
		rgb:  0xFFFFFF00,
		hsl:  0xFF27FF7F,
	},
	{
		desc: "green",
		rgb:  0xFF00FF00,
		hsl:  0xFF55FF7F,
	},
	{
		desc: "cyan",
		rgb:  0xFF00FFFF,
		hsl:  0xFF7FFF7F,
	},
	{
		desc: "blue",
		rgb:  0xFF0000FF,
		hsl:  0xFFAAFF7F,
	},
	{
		desc: "magenta",
		rgb:  0xFFFF00FF,
		hsl:  0xFFD9FF7F,
	},

	{
		desc: "white",
		rgb:  0xFFFFFFFF,
		hsl:  0xFF0000FF,
	},
}

const hue uint32 = 0xFF0000
const sat uint32 = 0xFF00 // saturation
const lum uint32 = 0xFF   // luma

func TestRGBtoHSL(t *testing.T) {
	for _, test := range colors {
		hslFromCode := RGBtoHSL(test.rgb)
		assert.Equal(t, test.hsl&hue, hslFromCode&hue, test.desc+"\t: Hue")
		assert.Equal(t, test.hsl&sat, hslFromCode&sat, test.desc+"\t: Saturation")
		assert.Equal(t, test.hsl&lum, hslFromCode&lum, test.desc+"\t: Luminance")
//		assert.Equal(t, test.hsl, hslFromCode, test.desc+"\t: All")
	}
}

func TestHSLtoRGB(t *testing.T) {
	for _, test := range colors {
		rgbFromCode := HSLtoRGB(test.hsl)
		assert.Equal(t, test.rgb&hue, rgbFromCode&hue, test.desc+"\t: Hue")
		assert.Equal(t, test.rgb&sat, rgbFromCode&sat, test.desc+"\t: Saturation")
		assert.Equal(t, test.rgb&lum, rgbFromCode&lum, test.desc+"\t: Luminance")
//		assert.Equal(t, test.rgb, rgbFromCode, test.desc+"\t: All")
	}
}
<commit_msg>Fixed HSLtoRGB test<commit_after>package main

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

var colors = []struct {
	desc string
	rgb  uint32
	hsl  uint32
}{
	{
		desc: "black",
		rgb:  0xFF000000,
		hsl:  0xFF000000,
	},
	{
		desc: "red",
		rgb:  0xFFFF0000,
		hsl:  0xFF00FF7F,
	},
	{
		desc: "yellow",
		rgb:  0xFFFFFF00,
		hsl:  0xFF27FF7F,
	},
	{
		desc: "green",
		rgb:  0xFF00FF00,
		hsl:  0xFF55FF7F,
	},
	{
		desc: "cyan",
		rgb:  0xFF00FFFF,
		hsl:  0xFF7FFF7F,
	},
	{
		desc: "blue",
		rgb:  0xFF0000FF,
		hsl:  0xFFAAFF7F,
	},
	{
		desc: "magenta",
		rgb:  0xFFFF00FF,
		hsl:  0xFFD9FF7F,
	},

	{
		desc: "white",
		rgb:  0xFFFFFFFF,
		hsl:  0xFF0000FF,
	},
}

const hue uint32 = 0xFF0000
const sat uint32 = 0xFF00 // saturation
const lum uint32 = 0xFF   // luma

func TestRGBtoHSL(t *testing.T) {
	for _, test := range colors {
		hslFromCode := RGBtoHSL(test.rgb)
		assert.Equal(t, test.hsl&hue, hslFromCode&hue, test.desc+"\t: Hue")
		assert.Equal(t, test.hsl&sat, hslFromCode&sat, test.desc+"\t: Saturation")
		assert.Equal(t, test.hsl&lum, hslFromCode&lum, test.desc+"\t: Luminance")
//		assert.Equal(t, test.hsl, hslFromCode, test.desc+"\t: All")
	}
}

func TestHSLtoRGB(t *testing.T) {
	for _, test := range colors {
		rgbFromCode := HSLtoRGB(test.hsl)
//		assert.Equal(t, test.rgb&hue, rgbFromCode&hue, test.desc+"\t: Hue")
//		assert.Equal(t, test.rgb&sat, rgbFromCode&sat, test.desc+"\t: Saturation")
//		assert.Equal(t, test.rgb&lum, rgbFromCode&lum, test.desc+"\t: Luminance")
		assert.Equal(t, test.rgb, rgbFromCode, test.desc+"\t: All")
	}
}
<|endoftext|>


<commit_before>package kontrol

import (
	"database/sql"
	"errors"
	"fmt"
	"net/url"
	"os"
	"regexp"
	"strings"
	"time"

	"github.com/hashicorp/go-version"
	_ "github.com/lib/pq"

	"github.com/koding/kite"
	kontrolprotocol "github.com/koding/kite/kontrol/protocol"
	"github.com/koding/kite/protocol"
)

// Postgres holds Postgresql database related configuration
type PostgresConfig struct {
	Host     string
	Port     int
	Username string
	Password string
	DBName   string
}

type Postgres struct {
	DB  *sql.DB
	Log kite.Logger
}

func NewPostgres(conf *PostgresConfig, log kite.Logger) *Postgres {
	if conf == nil {
		conf = &PostgresConfig{}
	}

	if conf.Port == 0 {
		conf.Port = 5432
	}

	if conf.Host == "" {
		conf.Host = "localhost"
	}

	if conf.DBName == "" {
		conf.DBName = os.Getenv("KONTROL_POSTGRES_DBNAME")
		if conf.DBName == "" {
			panic("db name is not set for postgres kontrol storage")
		}
	}

	connString := fmt.Sprintf(
		"host=%s port=%d dbname=%s sslmode=disable",
		conf.Host, conf.Port, conf.DBName,
	)

	if conf.Password != "" {
		connString += " password=" + conf.Password
	}

	if conf.Username == "" {
		conf.Username = os.Getenv("KONTROL_POSTGRES_USERNAME")
		if conf.Username == "" {
			panic("username is not set for postgres kontrol storage")
		}
	}

	connString += " user=" + conf.Username

	db, err := sql.Open("postgres", connString)
	if err != nil {
		panic(err)
	}

	// enable the ltree module which we are going to use, any error means it's
	// failed so there is no sense to continue, panic!
	enableTree := `CREATE EXTENSION IF NOT EXISTS ltree`
	if _, err := db.Exec(enableTree); err != nil {
		panic(err)
	}

	// create our initial kites table
	// * kite is going to be our ltree
	// * url is containing the kite's register url
	// * id is going to be kites' unique id (which also exists in the ltree
	//   path). We are adding it as a primary key so each kite with the full
	//   path can only exist once.
	// * created_at and updated_at are updated at creation and updating (like
	//  if the URL has changed)
	// Some notes:
	// * path label can only contain a sequence of alphanumeric characters
	//   and underscores. So for example a version string of "1.0.4" needs to
	//   be converted to "1_0_4" or uuid of 1111-2222-3333-4444 needs to be
	//   converted to 1111_2222_3333_4444.
	table := `CREATE TABLE IF NOT EXISTS kites (
		kite ltree NOT NULL,
		url text NOT NULL,
		id uuid PRIMARY KEY,
		created_at timestamptz NOT NULL DEFAULT (NOW() AT TIME ZONE 'UTC'),
		updated_at timestamptz NOT NULL DEFAULT (NOW() AT TIME ZONE 'UTC')
	);`

	if _, err := db.Exec(table); err != nil {
		panic(err)
	}

	// We enable index on the kite field. We don't return on errors because the
	// operator `IF NOT EXISTS` doesn't work for index creation, therefore we
	// assume the indexes might be already created.
	enableGistIndex := `CREATE INDEX kite_path_gist_idx ON kites USING GIST(kite)`
	enableBtreeIndex := `CREATE INDEX kite_path_btree_idx ON kites USING BTREE(kite)`

	if _, err := db.Exec(enableGistIndex); err != nil {
		log.Warning("postgres: enable gist index: %s", err)
	}

	if _, err := db.Exec(enableBtreeIndex); err != nil {
		log.Warning("postgres: enable btree index: %s", err)
	}

	p := &Postgres{
		DB:  db,
		Log: log,
	}

	cleanInterval := 30 * time.Second  // clean every 30 second
	expireInterval := 10 * time.Second // clean rows that are 10 second old
	go p.RunCleaner(cleanInterval, expireInterval)

	return p
}

// RunCleaner delets every "interval" duration rows which are older than
// "expire" duration based on the "updated_at" field. For more info check
// CleanExpireRows which is used to delete old rows.
func (p *Postgres) RunCleaner(interval, expire time.Duration) {
	cleanFunc := func() {
		affectedRows, err := p.CleanExpiredRows(expire)
		if err != nil {
			p.Log.Warning("postgres: cleaning old rows failed: %s", err)
		} else if affectedRows != 0 {
			p.Log.Info("postgres: cleaned up %d rows", affectedRows)
		}
	}

	cleanFunc() // run for the first time
	for _ = range time.Tick(interval) {
		cleanFunc()
	}
}

// CleanExpiredRows deletes rows that are at least "expire" duration old. So if
// say an expire duration of 10 second is given, it will delete all rows that
// were updated 10 seconds ago
func (p *Postgres) CleanExpiredRows(expire time.Duration) (int64, error) {
	// See: http://stackoverflow.com/questions/14465727/how-to-insert-things-like-now-interval-2-minutes-into-php-pdo-query
	// basically by passing an integer to INTERVAL is not possible, we need to
	// cast it. However there is a more simpler way, we can multiply INTERVAL
	// with an integer so we just declare a one second INTERVAL and multiply it
	// with the amount we want.
	cleanOldRows := `DELETE FROM kites WHERE updated_at < (now() at time zone 'utc') - ((INTERVAL '1 second') * $1)`

	rows, err := p.DB.Exec(cleanOldRows, int64(expire/time.Second))
	if err != nil {
		return 0, err
	}

	return rows.RowsAffected()
}

func (p *Postgres) Get(query *protocol.KontrolQuery) (Kites, error) {
	// only let query with usernames, otherwise the whole tree will be fetched
	// which is not good for us
	if query.Username == "" {
		return nil, errors.New("username is not specified in query")
	}

	path := ltreePath(query)

	var hasVersionConstraint bool // does query contains a constraint on version?
	var keyRest string            // query key after the version field
	var versionConstraint version.Constraints
	// NewVersion returns an error if it's a constraint, like: ">= 1.0, < 1.4"
	_, err := version.NewVersion(query.Version)
	if err != nil && query.Version != "" {
		// now parse our constraint
		versionConstraint, err = version.NewConstraint(query.Version)
		if err != nil {
			// version is a malformed, just return the error
			return nil, err
		}

		hasVersionConstraint = true
		nameQuery := &protocol.KontrolQuery{
			Username:    query.Username,
			Environment: query.Environment,
			Name:        query.Name,
		}

		// We will make a get request to all nodes under this name
		// and filter the result later.
		path = ltreePath(nameQuery)

		// Rest of the key after version field
		keyRest = "/" + strings.TrimRight(
			query.Region+"/"+query.Hostname+"/"+query.ID, "/")

	}

	rows, err := p.DB.Query(`SELECT kite, url FROM kites WHERE kite <@ $1`, path)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var kitePath string
	var url string

	kites := make(Kites, 0)

	for rows.Next() {
		err := rows.Scan(&kitePath, &url)
		if err != nil {
			return nil, err
		}

		kiteProt, err := kiteFromPath(kitePath)
		if err != nil {
			return nil, err
		}

		kites = append(kites, &protocol.KiteWithToken{
			Kite: *kiteProt,
			URL:  url,
		})
	}

	if err := rows.Err(); err != nil {
		return nil, err
	}

	// if it's just single result there is no need to shuffle or filter
	// according to the version constraint
	if len(kites) == 1 {
		return kites, nil
	}

	// Filter kites by version constraint
	if hasVersionConstraint {
		kites.Filter(versionConstraint, keyRest)
	}

	// randomize the result
	kites.Shuffle()

	return kites, nil
}

func (p *Postgres) Add(kiteProt *protocol.Kite, value *kontrolprotocol.RegisterValue) error {
	// check that the incoming URL is valid to prevent malformed input
	_, err := url.Parse(value.URL)
	if err != nil {
		return err
	}

	_, err = p.DB.Exec("INSERT into kites(kite, url, id) VALUES($1, $2, $3)",
		ltreePath(kiteProt.Query()),
		value.URL,
		kiteProt.ID,
	)
	return err
}

func (p *Postgres) Update(kiteProt *protocol.Kite, value *kontrolprotocol.RegisterValue) error {
	// check that the incoming url is valid to prevent malformed input
	_, err := url.Parse(value.URL)
	if err != nil {
		return err
	}

	// TODO: also consider just using WHERE id = kiteProt.ID, see how it's
	// performs out
	_, err = p.DB.Exec(`UPDATE kites SET url = $1, updated_at = (now() at time zone 'utc')
	WHERE kite ~ $2`,
		value.URL, ltreePath(kiteProt.Query()))

	return err
}

func (p *Postgres) Delete(kiteProt *protocol.Kite) error {
	deleteKite := `DELETE FROM kites WHERE kite ~ $1`
	_, err := p.DB.Exec(deleteKite, ltreePath(kiteProt.Query()))
	return err
}

func (p *Postgres) Clear() error {
	_, err := p.DB.Exec(`DROP TABLE kites`)
	return err
}

// ltreeLabel satisfies a invalid ltree definition of a label in path.
// According to the definition it is: "A label is a sequence of alphanumeric
// characters and underscores (for example, in C locale the characters
// A-Za-z0-9_ are allowed). Labels must be less than 256 bytes long."
//
// We could express one character with "[A-Za-z0-9_]", a word with
// "[A-Za-z0-9_]+". However we want to catch words that are not valid labels so
// we negate them with the "^" character, so it will be : "[^[A-Za-z0-9_]]+".
// Finally we cann use the POSIX character class: [:word:] which is:
// "Alphanumeric characters plus "_"", so the final regexp will be
// "[^[:word]]+"
var invalidLabelRe = regexp.MustCompile("[^[:word:]]+")

// ltreePath returns a query path to be used with the ltree module in postgress
// in the form of "username.environment.kitename.version.region.hostname.id"
func ltreePath(query *protocol.KontrolQuery) string {
	path := ""
	fields := query.Fields()

	// we stop for the first empty value
	for _, key := range keyOrder {
		v := fields[key]
		if v == "" {
			break
		}

		// replace anything that doesn't match the definition for a ltree path
		// label with a underscore, so the version "0.0.1" will be "0_0_1", or
		// uuid of "1111-2222-3333-4444" will be converted to
		// 1111_2222_3333_4444. Strings that satisfies the requirement are
		// untouched.
		v = invalidLabelRe.ReplaceAllLiteralString(v, "_")

		path = path + v + "."
	}

	// remove the latest dot which causes an invalid query
	path = strings.TrimSuffix(path, ".")
	return path
}

// kiteFromPath returns a Query from the given ltree path label
func kiteFromPath(path string) (*protocol.Kite, error) {
	fields := strings.Split(path, ".")

	if len(fields) != 7 {
		return nil, fmt.Errorf("invalid ltree path: %s", path)
	}

	// those labels were converted by us, therefore convert them back
	version := strings.Replace(fields[3], "_", ".", -1)
	id := strings.Replace(fields[6], "_", "-", -1)

	return &protocol.Kite{
		Username:    fields[0],
		Environment: fields[1],
		Name:        fields[2],
		Version:     version,
		Region:      fields[4],
		Hostname:    fields[5],
		ID:          id,
	}, nil

}
<commit_msg>kontrol/postgres: remove btree index because it doesn't use the operators we are using<commit_after>package kontrol

import (
	"database/sql"
	"errors"
	"fmt"
	"net/url"
	"os"
	"regexp"
	"strings"
	"time"

	"github.com/hashicorp/go-version"
	_ "github.com/lib/pq"

	"github.com/koding/kite"
	kontrolprotocol "github.com/koding/kite/kontrol/protocol"
	"github.com/koding/kite/protocol"
)

// Postgres holds Postgresql database related configuration
type PostgresConfig struct {
	Host     string
	Port     int
	Username string
	Password string
	DBName   string
}

type Postgres struct {
	DB  *sql.DB
	Log kite.Logger
}

func NewPostgres(conf *PostgresConfig, log kite.Logger) *Postgres {
	if conf == nil {
		conf = &PostgresConfig{}
	}

	if conf.Port == 0 {
		conf.Port = 5432
	}

	if conf.Host == "" {
		conf.Host = "localhost"
	}

	if conf.DBName == "" {
		conf.DBName = os.Getenv("KONTROL_POSTGRES_DBNAME")
		if conf.DBName == "" {
			panic("db name is not set for postgres kontrol storage")
		}
	}

	connString := fmt.Sprintf(
		"host=%s port=%d dbname=%s sslmode=disable",
		conf.Host, conf.Port, conf.DBName,
	)

	if conf.Password != "" {
		connString += " password=" + conf.Password
	}

	if conf.Username == "" {
		conf.Username = os.Getenv("KONTROL_POSTGRES_USERNAME")
		if conf.Username == "" {
			panic("username is not set for postgres kontrol storage")
		}
	}

	connString += " user=" + conf.Username

	db, err := sql.Open("postgres", connString)
	if err != nil {
		panic(err)
	}

	// enable the ltree module which we are going to use, any error means it's
	// failed so there is no sense to continue, panic!
	enableTree := `CREATE EXTENSION IF NOT EXISTS ltree`
	if _, err := db.Exec(enableTree); err != nil {
		panic(err)
	}

	// create our initial kites table
	// * kite is going to be our ltree
	// * url is containing the kite's register url
	// * id is going to be kites' unique id (which also exists in the ltree
	//   path). We are adding it as a primary key so each kite with the full
	//   path can only exist once.
	// * created_at and updated_at are updated at creation and updating (like
	//  if the URL has changed)
	// Some notes:
	// * path label can only contain a sequence of alphanumeric characters
	//   and underscores. So for example a version string of "1.0.4" needs to
	//   be converted to "1_0_4" or uuid of 1111-2222-3333-4444 needs to be
	//   converted to 1111_2222_3333_4444.
	table := `CREATE TABLE IF NOT EXISTS kites (
		kite ltree NOT NULL,
		url text NOT NULL,
		id uuid PRIMARY KEY,
		created_at timestamptz NOT NULL DEFAULT (NOW() AT TIME ZONE 'UTC'),
		updated_at timestamptz NOT NULL DEFAULT (NOW() AT TIME ZONE 'UTC')
	);`

	if _, err := db.Exec(table); err != nil {
		panic(err)
	}

	// We enable index on the kite field. We don't return on errors because the
	// operator `IF NOT EXISTS` doesn't work for index creation, therefore we
	// assume the indexes might be already created.
	enableGistIndex := `CREATE INDEX kite_path_gist_idx ON kites USING GIST(kite)`
	if _, err := db.Exec(enableGistIndex); err != nil {
		log.Warning("postgres: enable gist index: %s", err)
	}

	p := &Postgres{
		DB:  db,
		Log: log,
	}

	cleanInterval := 30 * time.Second  // clean every 30 second
	expireInterval := 10 * time.Second // clean rows that are 10 second old
	go p.RunCleaner(cleanInterval, expireInterval)

	return p
}

// RunCleaner delets every "interval" duration rows which are older than
// "expire" duration based on the "updated_at" field. For more info check
// CleanExpireRows which is used to delete old rows.
func (p *Postgres) RunCleaner(interval, expire time.Duration) {
	cleanFunc := func() {
		affectedRows, err := p.CleanExpiredRows(expire)
		if err != nil {
			p.Log.Warning("postgres: cleaning old rows failed: %s", err)
		} else if affectedRows != 0 {
			p.Log.Info("postgres: cleaned up %d rows", affectedRows)
		}
	}

	cleanFunc() // run for the first time
	for _ = range time.Tick(interval) {
		cleanFunc()
	}
}

// CleanExpiredRows deletes rows that are at least "expire" duration old. So if
// say an expire duration of 10 second is given, it will delete all rows that
// were updated 10 seconds ago
func (p *Postgres) CleanExpiredRows(expire time.Duration) (int64, error) {
	// See: http://stackoverflow.com/questions/14465727/how-to-insert-things-like-now-interval-2-minutes-into-php-pdo-query
	// basically by passing an integer to INTERVAL is not possible, we need to
	// cast it. However there is a more simpler way, we can multiply INTERVAL
	// with an integer so we just declare a one second INTERVAL and multiply it
	// with the amount we want.
	cleanOldRows := `DELETE FROM kites WHERE updated_at < (now() at time zone 'utc') - ((INTERVAL '1 second') * $1)`

	rows, err := p.DB.Exec(cleanOldRows, int64(expire/time.Second))
	if err != nil {
		return 0, err
	}

	return rows.RowsAffected()
}

func (p *Postgres) Get(query *protocol.KontrolQuery) (Kites, error) {
	// only let query with usernames, otherwise the whole tree will be fetched
	// which is not good for us
	if query.Username == "" {
		return nil, errors.New("username is not specified in query")
	}

	path := ltreePath(query)

	var hasVersionConstraint bool // does query contains a constraint on version?
	var keyRest string            // query key after the version field
	var versionConstraint version.Constraints
	// NewVersion returns an error if it's a constraint, like: ">= 1.0, < 1.4"
	_, err := version.NewVersion(query.Version)
	if err != nil && query.Version != "" {
		// now parse our constraint
		versionConstraint, err = version.NewConstraint(query.Version)
		if err != nil {
			// version is a malformed, just return the error
			return nil, err
		}

		hasVersionConstraint = true
		nameQuery := &protocol.KontrolQuery{
			Username:    query.Username,
			Environment: query.Environment,
			Name:        query.Name,
		}

		// We will make a get request to all nodes under this name
		// and filter the result later.
		path = ltreePath(nameQuery)

		// Rest of the key after version field
		keyRest = "/" + strings.TrimRight(
			query.Region+"/"+query.Hostname+"/"+query.ID, "/")

	}

	rows, err := p.DB.Query(`SELECT kite, url FROM kites WHERE kite <@ $1`, path)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var kitePath string
	var url string

	kites := make(Kites, 0)

	for rows.Next() {
		err := rows.Scan(&kitePath, &url)
		if err != nil {
			return nil, err
		}

		kiteProt, err := kiteFromPath(kitePath)
		if err != nil {
			return nil, err
		}

		kites = append(kites, &protocol.KiteWithToken{
			Kite: *kiteProt,
			URL:  url,
		})
	}

	if err := rows.Err(); err != nil {
		return nil, err
	}

	// if it's just single result there is no need to shuffle or filter
	// according to the version constraint
	if len(kites) == 1 {
		return kites, nil
	}

	// Filter kites by version constraint
	if hasVersionConstraint {
		kites.Filter(versionConstraint, keyRest)
	}

	// randomize the result
	kites.Shuffle()

	return kites, nil
}

func (p *Postgres) Add(kiteProt *protocol.Kite, value *kontrolprotocol.RegisterValue) error {
	// check that the incoming URL is valid to prevent malformed input
	_, err := url.Parse(value.URL)
	if err != nil {
		return err
	}

	_, err = p.DB.Exec("INSERT into kites(kite, url, id) VALUES($1, $2, $3)",
		ltreePath(kiteProt.Query()),
		value.URL,
		kiteProt.ID,
	)
	return err
}

func (p *Postgres) Update(kiteProt *protocol.Kite, value *kontrolprotocol.RegisterValue) error {
	// check that the incoming url is valid to prevent malformed input
	_, err := url.Parse(value.URL)
	if err != nil {
		return err
	}

	// TODO: also consider just using WHERE id = kiteProt.ID, see how it's
	// performs out
	_, err = p.DB.Exec(`UPDATE kites SET url = $1, updated_at = (now() at time zone 'utc')
	WHERE kite ~ $2`,
		value.URL, ltreePath(kiteProt.Query()))

	return err
}

func (p *Postgres) Delete(kiteProt *protocol.Kite) error {
	deleteKite := `DELETE FROM kites WHERE kite ~ $1`
	_, err := p.DB.Exec(deleteKite, ltreePath(kiteProt.Query()))
	return err
}

func (p *Postgres) Clear() error {
	_, err := p.DB.Exec(`DROP TABLE kites`)
	return err
}

// ltreeLabel satisfies a invalid ltree definition of a label in path.
// According to the definition it is: "A label is a sequence of alphanumeric
// characters and underscores (for example, in C locale the characters
// A-Za-z0-9_ are allowed). Labels must be less than 256 bytes long."
//
// We could express one character with "[A-Za-z0-9_]", a word with
// "[A-Za-z0-9_]+". However we want to catch words that are not valid labels so
// we negate them with the "^" character, so it will be : "[^[A-Za-z0-9_]]+".
// Finally we cann use the POSIX character class: [:word:] which is:
// "Alphanumeric characters plus "_"", so the final regexp will be
// "[^[:word]]+"
var invalidLabelRe = regexp.MustCompile("[^[:word:]]+")

// ltreePath returns a query path to be used with the ltree module in postgress
// in the form of "username.environment.kitename.version.region.hostname.id"
func ltreePath(query *protocol.KontrolQuery) string {
	path := ""
	fields := query.Fields()

	// we stop for the first empty value
	for _, key := range keyOrder {
		v := fields[key]
		if v == "" {
			break
		}

		// replace anything that doesn't match the definition for a ltree path
		// label with a underscore, so the version "0.0.1" will be "0_0_1", or
		// uuid of "1111-2222-3333-4444" will be converted to
		// 1111_2222_3333_4444. Strings that satisfies the requirement are
		// untouched.
		v = invalidLabelRe.ReplaceAllLiteralString(v, "_")

		path = path + v + "."
	}

	// remove the latest dot which causes an invalid query
	path = strings.TrimSuffix(path, ".")
	return path
}

// kiteFromPath returns a Query from the given ltree path label
func kiteFromPath(path string) (*protocol.Kite, error) {
	fields := strings.Split(path, ".")

	if len(fields) != 7 {
		return nil, fmt.Errorf("invalid ltree path: %s", path)
	}

	// those labels were converted by us, therefore convert them back
	version := strings.Replace(fields[3], "_", ".", -1)
	id := strings.Replace(fields[6], "_", "-", -1)

	return &protocol.Kite{
		Username:    fields[0],
		Environment: fields[1],
		Name:        fields[2],
		Version:     version,
		Region:      fields[4],
		Hostname:    fields[5],
		ID:          id,
	}, nil

}
<|endoftext|>


<commit_before>package shadowsocks

import (
	"bytes"
	"fmt"
	"io"
	"strings"
	"testing"

	"github.com/shadowsocks/go-shadowsocks2/shadowaead"
	"golang.org/x/crypto/chacha20poly1305"
)

func newTestCipher(t *testing.T) shadowaead.Cipher {
	key := []byte("12345678901234567890123456789012") // 32 bytes
	cipher, err := shadowaead.Chacha20Poly1305(key)
	if err != nil {
		t.Fatal(err)
	}
	return cipher
}

// Overhead for cipher chacha20poly1305
const testCipherOverhead = 16

func TestCipherReaderAuthentiationFailure(t *testing.T) {
	cipher := newTestCipher(t)

	clientReader := strings.NewReader("Fails Authentication")
	reader := NewShadowsocksReader(clientReader, cipher)
	_, err := reader.Read(make([]byte, 1))
	if err == nil {
		t.Fatalf("Expected authentication failure, got %v", err)
	}
}

func TestCipherReaderUnexpectedEOF(t *testing.T) {
	cipher := newTestCipher(t)

	clientReader := strings.NewReader("short")
	server := NewShadowsocksReader(clientReader, cipher)
	_, err := server.Read(make([]byte, 10))
	if err != io.ErrUnexpectedEOF {
		t.Fatalf("Expected ErrUnexpectedEOF, got %v", err)
	}
}

func TestCipherReaderEOF(t *testing.T) {
	cipher := newTestCipher(t)

	clientReader := strings.NewReader("")
	server := NewShadowsocksReader(clientReader, cipher)
	_, err := server.Read(make([]byte, 10))
	if err != io.EOF {
		t.Fatalf("Expected EOF, got %v", err)
	}
	_, err = server.Read([]byte{})
	if err != io.EOF {
		t.Fatalf("Expected EOF, got %v", err)
	}
}

func encryptBlocks(cipher shadowaead.Cipher, salt []byte, blocks [][]byte) (io.Reader, error) {
	var ssText bytes.Buffer
	aead, err := cipher.Encrypter(salt)
	if err != nil {
		return nil, fmt.Errorf("Failed to create AEAD: %v", err)
	}
	ssText.Write(salt)
	// buf must fit the larges block ciphertext
	buf := make([]byte, 2+100+testCipherOverhead)
	var expectedCipherSize int
	nonce := make([]byte, chacha20poly1305.NonceSize)
	for _, block := range blocks {
		ssText.Write(aead.Seal(buf[:0], nonce, []byte{0, byte(len(block))}, nil))
		nonce[0]++
		expectedCipherSize += 2 + testCipherOverhead
		ssText.Write(aead.Seal(buf[:0], nonce, block, nil))
		nonce[0]++
		expectedCipherSize += len(block) + testCipherOverhead
	}
	if ssText.Len() != cipher.SaltSize()+expectedCipherSize {
		return nil, fmt.Errorf("cipherText has size %v. Expected %v", ssText.Len(), cipher.SaltSize()+expectedCipherSize)
	}
	return &ssText, nil
}

func TestCipherReaderGoodReads(t *testing.T) {
	cipher := newTestCipher(t)

	salt := []byte("12345678901234567890123456789012")
	if len(salt) != cipher.SaltSize() {
		t.Fatalf("Salt has size %v. Expected %v", len(salt), cipher.SaltSize())
	}
	ssText, err := encryptBlocks(cipher, salt, [][]byte{
		[]byte("[First Block]"),
		[]byte(""), // Corner case: empty block
		[]byte("[Third Block]")})
	if err != nil {
		t.Fatal(err)
	}

	reader := NewShadowsocksReader(ssText, cipher)
	plainText := make([]byte, len("[First Block]")+len("[Third Block]"))
	n, err := io.ReadFull(reader, plainText)
	if err != nil {
		t.Fatalf("Failed to fully read plain text. Got %v bytes: %v", n, err)
	}
	_, err = reader.Read([]byte{})
	if err != io.EOF {
		t.Fatalf("Expected EOF, got %v", err)
	}
	_, err = reader.Read(make([]byte, 1))
	if err != io.EOF {
		t.Fatalf("Expected EOF, got %v", err)
	}
}

func TestCipherReaderClose(t *testing.T) {
	cipher := newTestCipher(t)

	pipeReader, pipeWriter := io.Pipe()
	server := NewShadowsocksReader(pipeReader, cipher)
	result := make(chan error)
	go func() {
		_, err := server.Read(make([]byte, 10))
		result <- err
	}()
	pipeWriter.Close()
	err := <-result
	if err != io.EOF {
		t.Fatalf("Expected ErrUnexpectedEOF, got %v", err)
	}
}

func TestCipherReaderCloseError(t *testing.T) {
	cipher := newTestCipher(t)

	pipeReader, pipeWriter := io.Pipe()
	server := NewShadowsocksReader(pipeReader, cipher)
	result := make(chan error)
	go func() {
		_, err := server.Read(make([]byte, 10))
		result <- err
	}()
	pipeWriter.CloseWithError(fmt.Errorf("xx!!ERROR!!xx"))
	err := <-result
	if err == nil || !strings.Contains(err.Error(), "xx!!ERROR!!xx") {
		t.Fatalf("Unexpected error: %v", err)
	}
}

func TestEndToEnd(t *testing.T) {
	cipher := newTestCipher(t)

	connReader, connWriter := io.Pipe()
	writer := NewShadowsocksWriter(connWriter, cipher)
	reader := NewShadowsocksReader(connReader, cipher)
	expected := "Test"
	go func() {
		defer connWriter.Close()
		_, err := writer.Write([]byte(expected))
		if err != nil {
			t.Fatalf("Failed Write: %v", err)
		}
	}()
	var output bytes.Buffer
	_, err := reader.WriteTo(&output)
	if err != nil {
		t.Fatalf("Failed WriteTo: %v", err)
	}
	if output.String() != expected {
		t.Fatalf("Expected output '%v'. Got '%v'", expected, output.String())
	}
}
<commit_msg>spelling: authentication<commit_after>package shadowsocks

import (
	"bytes"
	"fmt"
	"io"
	"strings"
	"testing"

	"github.com/shadowsocks/go-shadowsocks2/shadowaead"
	"golang.org/x/crypto/chacha20poly1305"
)

func newTestCipher(t *testing.T) shadowaead.Cipher {
	key := []byte("12345678901234567890123456789012") // 32 bytes
	cipher, err := shadowaead.Chacha20Poly1305(key)
	if err != nil {
		t.Fatal(err)
	}
	return cipher
}

// Overhead for cipher chacha20poly1305
const testCipherOverhead = 16

func TestCipherReaderAuthenticationFailure(t *testing.T) {
	cipher := newTestCipher(t)

	clientReader := strings.NewReader("Fails Authentication")
	reader := NewShadowsocksReader(clientReader, cipher)
	_, err := reader.Read(make([]byte, 1))
	if err == nil {
		t.Fatalf("Expected authentication failure, got %v", err)
	}
}

func TestCipherReaderUnexpectedEOF(t *testing.T) {
	cipher := newTestCipher(t)

	clientReader := strings.NewReader("short")
	server := NewShadowsocksReader(clientReader, cipher)
	_, err := server.Read(make([]byte, 10))
	if err != io.ErrUnexpectedEOF {
		t.Fatalf("Expected ErrUnexpectedEOF, got %v", err)
	}
}

func TestCipherReaderEOF(t *testing.T) {
	cipher := newTestCipher(t)

	clientReader := strings.NewReader("")
	server := NewShadowsocksReader(clientReader, cipher)
	_, err := server.Read(make([]byte, 10))
	if err != io.EOF {
		t.Fatalf("Expected EOF, got %v", err)
	}
	_, err = server.Read([]byte{})
	if err != io.EOF {
		t.Fatalf("Expected EOF, got %v", err)
	}
}

func encryptBlocks(cipher shadowaead.Cipher, salt []byte, blocks [][]byte) (io.Reader, error) {
	var ssText bytes.Buffer
	aead, err := cipher.Encrypter(salt)
	if err != nil {
		return nil, fmt.Errorf("Failed to create AEAD: %v", err)
	}
	ssText.Write(salt)
	// buf must fit the larges block ciphertext
	buf := make([]byte, 2+100+testCipherOverhead)
	var expectedCipherSize int
	nonce := make([]byte, chacha20poly1305.NonceSize)
	for _, block := range blocks {
		ssText.Write(aead.Seal(buf[:0], nonce, []byte{0, byte(len(block))}, nil))
		nonce[0]++
		expectedCipherSize += 2 + testCipherOverhead
		ssText.Write(aead.Seal(buf[:0], nonce, block, nil))
		nonce[0]++
		expectedCipherSize += len(block) + testCipherOverhead
	}
	if ssText.Len() != cipher.SaltSize()+expectedCipherSize {
		return nil, fmt.Errorf("cipherText has size %v. Expected %v", ssText.Len(), cipher.SaltSize()+expectedCipherSize)
	}
	return &ssText, nil
}

func TestCipherReaderGoodReads(t *testing.T) {
	cipher := newTestCipher(t)

	salt := []byte("12345678901234567890123456789012")
	if len(salt) != cipher.SaltSize() {
		t.Fatalf("Salt has size %v. Expected %v", len(salt), cipher.SaltSize())
	}
	ssText, err := encryptBlocks(cipher, salt, [][]byte{
		[]byte("[First Block]"),
		[]byte(""), // Corner case: empty block
		[]byte("[Third Block]")})
	if err != nil {
		t.Fatal(err)
	}

	reader := NewShadowsocksReader(ssText, cipher)
	plainText := make([]byte, len("[First Block]")+len("[Third Block]"))
	n, err := io.ReadFull(reader, plainText)
	if err != nil {
		t.Fatalf("Failed to fully read plain text. Got %v bytes: %v", n, err)
	}
	_, err = reader.Read([]byte{})
	if err != io.EOF {
		t.Fatalf("Expected EOF, got %v", err)
	}
	_, err = reader.Read(make([]byte, 1))
	if err != io.EOF {
		t.Fatalf("Expected EOF, got %v", err)
	}
}

func TestCipherReaderClose(t *testing.T) {
	cipher := newTestCipher(t)

	pipeReader, pipeWriter := io.Pipe()
	server := NewShadowsocksReader(pipeReader, cipher)
	result := make(chan error)
	go func() {
		_, err := server.Read(make([]byte, 10))
		result <- err
	}()
	pipeWriter.Close()
	err := <-result
	if err != io.EOF {
		t.Fatalf("Expected ErrUnexpectedEOF, got %v", err)
	}
}

func TestCipherReaderCloseError(t *testing.T) {
	cipher := newTestCipher(t)

	pipeReader, pipeWriter := io.Pipe()
	server := NewShadowsocksReader(pipeReader, cipher)
	result := make(chan error)
	go func() {
		_, err := server.Read(make([]byte, 10))
		result <- err
	}()
	pipeWriter.CloseWithError(fmt.Errorf("xx!!ERROR!!xx"))
	err := <-result
	if err == nil || !strings.Contains(err.Error(), "xx!!ERROR!!xx") {
		t.Fatalf("Unexpected error: %v", err)
	}
}

func TestEndToEnd(t *testing.T) {
	cipher := newTestCipher(t)

	connReader, connWriter := io.Pipe()
	writer := NewShadowsocksWriter(connWriter, cipher)
	reader := NewShadowsocksReader(connReader, cipher)
	expected := "Test"
	go func() {
		defer connWriter.Close()
		_, err := writer.Write([]byte(expected))
		if err != nil {
			t.Fatalf("Failed Write: %v", err)
		}
	}()
	var output bytes.Buffer
	_, err := reader.WriteTo(&output)
	if err != nil {
		t.Fatalf("Failed WriteTo: %v", err)
	}
	if output.String() != expected {
		t.Fatalf("Expected output '%v'. Got '%v'", expected, output.String())
	}
}
<|endoftext|>