<commit_before>package pqtgo_test

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

	"github.com/aryann/difflib"
	"github.com/piotrkowalczuk/pqt"
	"github.com/piotrkowalczuk/pqt/pqtgo"
)

func TestGenerator_Generate(t *testing.T) {
	cases := map[string]struct {
		schema    *pqt.Schema
		generator *pqtgo.Generator
		expected  string
	}{
		"basic": {
			schema: pqt.NewSchema("text"),
			generator: pqtgo.NewGenerator().
				AddImport("github.com/piotrkowalczuk/ntypes"),
			expected: `package main
import (
"github.com/go-kit/kit/log"
"github.com/m4rw3r/uuid"
"github.com/piotrkowalczuk/ntypes"
)
`,
		},
		"custom-package": {
			schema: pqt.NewSchema("text"),
			generator: pqtgo.NewGenerator().
				SetPackage("example").
				AddImport("fmt"),
			expected: `package example
import (
"github.com/go-kit/kit/log"
"github.com/m4rw3r/uuid"
"fmt"
)
`,
		},
		"simple table": {
			schema: pqt.NewSchema("text").AddTable(
				pqt.NewTable("first").AddColumn(
					pqt.NewColumn("id", pqt.TypeSerialBig()),
				).AddColumn(
					pqt.NewColumn("name", pqt.TypeText()),
				),
			),
			generator: pqtgo.NewGenerator().
				SetPackage("custom"),
			expected: `package custom
import (
"github.com/go-kit/kit/log"
"github.com/m4rw3r/uuid"
)
const (
tableFirst = "text.first"
	tableFirstColumnId = "id"
		tableFirstColumnName = "name"
		)
var (
tableFirstColumns = []string{
tableFirstColumnId,
tableFirstColumnName,
})
type firstEntity struct{
// id ...
id *ntypes.Int64
// name ...
name *ntypes.String
}

func (e *firstEntity) prop(cn string) (interface{}, bool) {
switch cn {
case tableFirstColumnId:
return &e.id, true
case tableFirstColumnName:
return &e.name, true
default:
return nil, false
}
}
func (e *firstEntity) props(cns ...string) ([]interface{}, error) {

		res := make([]interface{}, 0, len(cns))
		for _, cn := range cns {
			if prop, ok := e.prop(cn); ok {
				res = append(res, prop)
			} else {
				return nil, fmt.Errorf("unexpected column provided: %s", cn)
			}
		}
		return res, nil
}

// firstIterator is not thread safe.
type firstIterator struct {
	rows *sql.Rows
	cols []string
}

func (i *firstIterator) Next() bool {
	return i.rows.Next()
}

func (i *firstIterator) Close() error {
	return i.rows.Close()
}

func (i *firstIterator) Err() error {
	return i.rows.Err()
}

// Columns is wrapper around sql.Rows.Columns method, that also cache outpu inside iterator.
func (i *firstIterator) Columns() ([]string, error) {
	if i.cols == nil {
		cols, err := i.rows.Columns()
		if err != nil {
			return nil, err
		}
		i.cols = cols
	}
	return i.cols, nil
}

// Ent is wrapper around first method that makes iterator more generic.
func (i *firstIterator) Ent() (interface{}, error) {
	return i.First()
}

func (i *firstIterator) First() (*firstEntity, error) {
	var ent firstEntity
	cols, err := i.rows.Columns()
	if err != nil {
		return nil, err
	}

	props, err := ent.props(cols...)
	if err != nil {
		return nil, err
	}
	if err := i.rows.Scan(props...); err != nil {
		return nil, err
	}
	return &ent, nil
}
type firstCriteria struct {
offset, limit int64
sort map[string]bool
id *qtypes.Int64
name *qtypes.String
}

func (c *firstCriteria) WriteComposition(sel string, com *pqtgo.Composer, opt *pqtgo.CompositionOpts) (err error) {

		if err = pqtgo.WriteCompositionQueryInt64(c.id, tableFirstColumnId, com, pqtgo.And); err != nil {
			return
		}

		if err = pqtgo.WriteCompositionQueryString(c.name, tableFirstColumnName, com, pqtgo.And); err != nil {
			return
		}

	if len(c.sort) > 0 {
		i:=0
		com.WriteString(" ORDER BY ")

		for cn, asc := range c.sort {
			for _, tcn := range tableFirstColumns {
				if cn == tcn {
					if i > 0 {
						com.WriteString(", ")
					}
					com.WriteString(cn)
					if !asc {
						com.WriteString(" DESC ")
					}
					i++
					break
				}
			}
		}
	}
	if c.offset > 0 {
		if _, err = com.WriteString(" OFFSET "); err != nil {
			return
		}
		if err = com.WritePlaceholder(); err != nil {
			return
		}
		if _, err = com.WriteString(" "); err != nil {
			return
		}
		com.Add(c.offset)
	}
	if c.limit > 0 {
		if _, err = com.WriteString(" LIMIT "); err != nil {
			return
		}
		if err = com.WritePlaceholder(); err != nil {
			return
		}
		if _, err = com.WriteString(" "); err != nil {
			return
		}
		com.Add(c.limit)
	}

	return
}
type firstPatch struct {
id *ntypes.Int64
name *ntypes.String
}

		type firstRepositoryBase struct {
			table string
			columns []string
			db *sql.DB
			dbg bool
			log log.Logger
		}
	func scanFirstRows(rows *sql.Rows) ([]*firstEntity, error) {
	var (
		entities []*firstEntity
		err error
	)
	for rows.Next() {
		var ent firstEntity
		err = rows.Scan(
	&ent.id,
&ent.name,
)
			if err != nil {
				return nil, err
			}

			entities = append(entities, &ent)
		}
		if rows.Err() != nil {
			return nil, rows.Err()
		}

		return entities, nil
	}

	func (r *firstRepositoryBase) count(c *firstCriteria) (int64, error) {

	com := pqtgo.NewComposer(2)
	buf := bytes.NewBufferString("SELECT COUNT(*) FROM ")
	buf.WriteString(r.table)

	if err := c.WriteComposition("", com, pqtgo.And); err != nil {
		return 0, err
	}
	if com.Dirty {
		buf.WriteString(" WHERE ")
	}
	if com.Len() > 0 {
		buf.ReadFrom(com)
	}

	if r.dbg {
		if err := r.log.Log("msg", buf.String(), "function", "Count"); err != nil {
			return 0, err
		}
	}

	var count int64
	if err := r.db.QueryRow(buf.String(), com.Args()...).Scan(&count); err != nil {
		return 0, err
	}
	return count, nil
}

func (r *firstRepositoryBase) find(c *firstCriteria) ([]*firstEntity, error) {

	com := pqtgo.NewComposer(1)
	buf := bytes.NewBufferString("SELECT ")
	buf.WriteString(strings.Join(r.columns, ", "))
	buf.WriteString(" FROM ")
	buf.WriteString(r.table)
	buf.WriteString(" ")

	if err := c.WriteComposition("", com, pqtgo.And); err != nil {
		return nil, err
	}
	if com.Dirty {
		buf.WriteString(" WHERE ")
	}
	if com.Len() > 0 {
		buf.ReadFrom(com)
	}

	if r.dbg {
		if err := r.log.Log("msg", buf.String(), "function", "Find"); err != nil {
			return nil, err
		}
	}

	rows, err := r.db.Query(buf.String(), com.Args()...)
	if err != nil {
		return nil, err
	}

	defer rows.Close()

	return scanFirstRows(rows)
}
func (r *firstRepositoryBase) findIter(c *firstCriteria) (*firstIterator, error) {

	com := pqtgo.NewComposer(1)
	buf := bytes.NewBufferString("SELECT ")
	buf.WriteString(strings.Join(r.columns, ", "))
	buf.WriteString(" FROM ")
	buf.WriteString(r.table)
	buf.WriteString(" ")

	if err := c.WriteComposition("", com, pqtgo.And); err != nil {
		return nil, err
	}
	if com.Dirty {
		buf.WriteString(" WHERE ")
	}
	if com.Len() > 0 {
		buf.ReadFrom(com)
	}

	if r.dbg {
		if err := r.log.Log("msg", buf.String(), "function", "Find"); err != nil {
			return nil, err
		}
	}

	rows, err := r.db.Query(buf.String(), com.Args()...)
	if err != nil {
		return nil, err
	}

	return &firstIterator{rows: rows}, nil
}
func (r *firstRepositoryBase) insert(e *firstEntity) (*firstEntity, error) {
		insert := pqcomp.New(0, 2)
	insert.AddExpr(tableFirstColumnName, "", e.name)

		b := bytes.NewBufferString("INSERT INTO " + r.table)

		if insert.Len() != 0 {
			b.WriteString(" (")
			for insert.Next() {
				if !insert.First() {
					b.WriteString(", ")
				}

				fmt.Fprintf(b, "%s", insert.Key())
			}
			insert.Reset()
			b.WriteString(") VALUES (")
			for insert.Next() {
				if !insert.First() {
					b.WriteString(", ")
				}

				fmt.Fprintf(b, "%s", insert.PlaceHolder())
			}
			b.WriteString(")")
			if len(r.columns) > 0 {
				b.WriteString(" RETURNING ")
				b.WriteString(strings.Join(r.columns, ", "))
			}
		}

		if r.dbg {
			if err := r.log.Log("msg", b.String(), "function", "Insert"); err != nil {
				return nil, err
			}
		}

		err := r.db.QueryRow(b.String(), insert.Args()...).Scan(
	&e.id,
&e.name,
)
		if err != nil {
			return nil, err
		}

		return e, nil
	}
func (r *firstRepositoryBase) upsert(e *firstEntity, p *firstPatch, inf ...string) (*firstEntity, error) {
		insert := pqcomp.New(0, 2)
		update := insert.Compose(2)
	insert.AddExpr(tableFirstColumnName, "", e.name)
if len(inf) > 0 {
update.AddExpr(tableFirstColumnName, "=", p.name)
}

		b := bytes.NewBufferString("INSERT INTO " + r.table)

		if insert.Len() > 0 {
			b.WriteString(" (")
			for insert.Next() {
				if !insert.First() {
					b.WriteString(", ")
				}

				fmt.Fprintf(b, "%s", insert.Key())
			}
			insert.Reset()
			b.WriteString(") VALUES (")
			for insert.Next() {
				if !insert.First() {
					b.WriteString(", ")
				}

				fmt.Fprintf(b, "%s", insert.PlaceHolder())
			}
			b.WriteString(")")
		}
		b.WriteString(" ON CONFLICT ")
		if len(inf) > 0 && update.Len() > 0 {
			b.WriteString(" (")
			for j, i := range inf {
				if j != 0 {
					b.WriteString(", ")
				}
				b.WriteString(i)
			}
			b.WriteString(") ")
			b.WriteString(" DO UPDATE SET ")
			for update.Next() {
				if !update.First() {
					b.WriteString(", ")
				}

				b.WriteString(update.Key())
				b.WriteString(" ")
				b.WriteString(update.Oper())
				b.WriteString(" ")
				b.WriteString(update.PlaceHolder())
			}
		} else {
			b.WriteString(" DO NOTHING ")
		}
		if insert.Len() > 0 {
			if len(r.columns) > 0 {
				b.WriteString(" RETURNING ")
				b.WriteString(strings.Join(r.columns, ", "))
			}
		}

		if r.dbg {
			if err := r.log.Log("msg", b.String(), "function", "Upsert"); err != nil {
				return nil, err
			}
		}

		err := r.db.QueryRow(b.String(), insert.Args()...).Scan(
	&e.id,
&e.name,
)
		if err != nil {
			return nil, err
		}

		return e, nil
	}
`,
		},
	}

	for hint, c := range cases {
		b, err := c.generator.Generate(c.schema)
		if err != nil {
			t.Errorf("%s: unexpected error: %s", hint, err.Error())
			continue
		}
		assertGoCode(t, c.expected, string(b), hint)
	}
}

func assertGoCode(t *testing.T, s1, s2, msg string, com ...interface{}) {
	s1 = fmt.Sprintf("%s", s1)
	s2 = fmt.Sprintf("%s", s2)
	tmp1 := strings.Split(s1, "\n")
	tmp2 := strings.Split(s2, "\n")
	if s1 != s2 {
		b := bytes.NewBuffer(nil)
		for _, diff := range difflib.Diff(tmp1, tmp2) {
			p := strings.Replace(diff.Payload, "\t", "\\t", -1)
			switch diff.Delta {
			case difflib.Common:
				fmt.Fprintf(b, "%s %s\n", diff.Delta.String(), p)
			case difflib.LeftOnly:
				fmt.Fprintf(b, "\033[31m%s %s\033[39m\n", diff.Delta.String(), p)
			case difflib.RightOnly:
				fmt.Fprintf(b, "\033[32m%s %s\033[39m\n", diff.Delta.String(), p)
			}
		}
		t.Errorf(b.String())
	}
}
<commit_msg>generator test fix<commit_after>package pqtgo_test

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

	"github.com/aryann/difflib"
	"github.com/piotrkowalczuk/pqt"
	"github.com/piotrkowalczuk/pqt/pqtgo"
)

func TestGenerator_Generate(t *testing.T) {
	cases := map[string]struct {
		schema    *pqt.Schema
		generator *pqtgo.Generator
		expected  string
	}{
		"basic": {
			schema: pqt.NewSchema("text"),
			generator: pqtgo.NewGenerator().
				AddImport("github.com/piotrkowalczuk/ntypes"),
			expected: `package main
import (
"github.com/go-kit/kit/log"
"github.com/m4rw3r/uuid"
"github.com/piotrkowalczuk/ntypes"
)
`,
		},
		"custom-package": {
			schema: pqt.NewSchema("text"),
			generator: pqtgo.NewGenerator().
				SetPackage("example").
				AddImport("fmt"),
			expected: `package example
import (
"github.com/go-kit/kit/log"
"github.com/m4rw3r/uuid"
"fmt"
)
`,
		},
		"simple table": {
			schema: pqt.NewSchema("text").AddTable(
				pqt.NewTable("first").AddColumn(
					pqt.NewColumn("id", pqt.TypeSerialBig()),
				).AddColumn(
					pqt.NewColumn("name", pqt.TypeText()),
				),
			),
			generator: pqtgo.NewGenerator().
				SetPackage("custom"),
			expected: `package custom
import (
"github.com/go-kit/kit/log"
"github.com/m4rw3r/uuid"
)
const (
tableFirst = "text.first"
	tableFirstColumnId = "id"
		tableFirstColumnName = "name"
		)
var (
tableFirstColumns = []string{
tableFirstColumnId,
tableFirstColumnName,
})
type firstEntity struct{
// id ...
id *ntypes.Int64
// name ...
name *ntypes.String
}

func (e *firstEntity) prop(cn string) (interface{}, bool) {
switch cn {
case tableFirstColumnId:
return &e.id, true
case tableFirstColumnName:
return &e.name, true
default:
return nil, false
}
}
func (e *firstEntity) props(cns ...string) ([]interface{}, error) {

		res := make([]interface{}, 0, len(cns))
		for _, cn := range cns {
			if prop, ok := e.prop(cn); ok {
				res = append(res, prop)
			} else {
				return nil, fmt.Errorf("unexpected column provided: %s", cn)
			}
		}
		return res, nil
}

// firstIterator is not thread safe.
type firstIterator struct {
	rows *sql.Rows
	cols []string
}

func (i *firstIterator) Next() bool {
	return i.rows.Next()
}

func (i *firstIterator) Close() error {
	return i.rows.Close()
}

func (i *firstIterator) Err() error {
	return i.rows.Err()
}

// Columns is wrapper around sql.Rows.Columns method, that also cache outpu inside iterator.
func (i *firstIterator) Columns() ([]string, error) {
	if i.cols == nil {
		cols, err := i.rows.Columns()
		if err != nil {
			return nil, err
		}
		i.cols = cols
	}
	return i.cols, nil
}

// Ent is wrapper around first method that makes iterator more generic.
func (i *firstIterator) Ent() (interface{}, error) {
	return i.First()
}

func (i *firstIterator) First() (*firstEntity, error) {
	var ent firstEntity
	cols, err := i.rows.Columns()
	if err != nil {
		return nil, err
	}

	props, err := ent.props(cols...)
	if err != nil {
		return nil, err
	}
	if err := i.rows.Scan(props...); err != nil {
		return nil, err
	}
	return &ent, nil
}
type firstCriteria struct {
offset, limit int64
sort map[string]bool
id *qtypes.Int64
name *qtypes.String
}

func (c *firstCriteria) WriteComposition(sel string, com *pqtgo.Composer, opt *pqtgo.CompositionOpts) (err error) {

		if err = pqtgo.WriteCompositionQueryInt64(c.id, tableFirstColumnId, com, &pqtgo.CompositionOpts{
		Joint: " AND ",
		IsJSON: false,
	}); err != nil {
			return
		}

		if err = pqtgo.WriteCompositionQueryString(c.name, tableFirstColumnName, com, pqtgo.And); err != nil {
			return
		}

	if len(c.sort) > 0 {
		i:=0
		com.WriteString(" ORDER BY ")

		for cn, asc := range c.sort {
			for _, tcn := range tableFirstColumns {
				if cn == tcn {
					if i > 0 {
						com.WriteString(", ")
					}
					com.WriteString(cn)
					if !asc {
						com.WriteString(" DESC ")
					}
					i++
					break
				}
			}
		}
	}
	if c.offset > 0 {
		if _, err = com.WriteString(" OFFSET "); err != nil {
			return
		}
		if err = com.WritePlaceholder(); err != nil {
			return
		}
		if _, err = com.WriteString(" "); err != nil {
			return
		}
		com.Add(c.offset)
	}
	if c.limit > 0 {
		if _, err = com.WriteString(" LIMIT "); err != nil {
			return
		}
		if err = com.WritePlaceholder(); err != nil {
			return
		}
		if _, err = com.WriteString(" "); err != nil {
			return
		}
		com.Add(c.limit)
	}

	return
}
type firstPatch struct {
id *ntypes.Int64
name *ntypes.String
}

		type firstRepositoryBase struct {
			table string
			columns []string
			db *sql.DB
			dbg bool
			log log.Logger
		}
	func scanFirstRows(rows *sql.Rows) ([]*firstEntity, error) {
	var (
		entities []*firstEntity
		err error
	)
	for rows.Next() {
		var ent firstEntity
		err = rows.Scan(
	&ent.id,
&ent.name,
)
			if err != nil {
				return nil, err
			}

			entities = append(entities, &ent)
		}
		if rows.Err() != nil {
			return nil, rows.Err()
		}

		return entities, nil
	}

	func (r *firstRepositoryBase) count(c *firstCriteria) (int64, error) {

	com := pqtgo.NewComposer(2)
	buf := bytes.NewBufferString("SELECT COUNT(*) FROM ")
	buf.WriteString(r.table)

	if err := c.WriteComposition("", com, pqtgo.And); err != nil {
		return 0, err
	}
	if com.Dirty {
		buf.WriteString(" WHERE ")
	}
	if com.Len() > 0 {
		buf.ReadFrom(com)
	}

	if r.dbg {
		if err := r.log.Log("msg", buf.String(), "function", "Count"); err != nil {
			return 0, err
		}
	}

	var count int64
	if err := r.db.QueryRow(buf.String(), com.Args()...).Scan(&count); err != nil {
		return 0, err
	}
	return count, nil
}

func (r *firstRepositoryBase) find(c *firstCriteria) ([]*firstEntity, error) {

	com := pqtgo.NewComposer(1)
	buf := bytes.NewBufferString("SELECT ")
	buf.WriteString(strings.Join(r.columns, ", "))
	buf.WriteString(" FROM ")
	buf.WriteString(r.table)
	buf.WriteString(" ")

	if err := c.WriteComposition("", com, pqtgo.And); err != nil {
		return nil, err
	}
	if com.Dirty {
		buf.WriteString(" WHERE ")
	}
	if com.Len() > 0 {
		buf.ReadFrom(com)
	}

	if r.dbg {
		if err := r.log.Log("msg", buf.String(), "function", "Find"); err != nil {
			return nil, err
		}
	}

	rows, err := r.db.Query(buf.String(), com.Args()...)
	if err != nil {
		return nil, err
	}

	defer rows.Close()

	return scanFirstRows(rows)
}
func (r *firstRepositoryBase) findIter(c *firstCriteria) (*firstIterator, error) {

	com := pqtgo.NewComposer(1)
	buf := bytes.NewBufferString("SELECT ")
	buf.WriteString(strings.Join(r.columns, ", "))
	buf.WriteString(" FROM ")
	buf.WriteString(r.table)
	buf.WriteString(" ")

	if err := c.WriteComposition("", com, pqtgo.And); err != nil {
		return nil, err
	}
	if com.Dirty {
		buf.WriteString(" WHERE ")
	}
	if com.Len() > 0 {
		buf.ReadFrom(com)
	}

	if r.dbg {
		if err := r.log.Log("msg", buf.String(), "function", "Find"); err != nil {
			return nil, err
		}
	}

	rows, err := r.db.Query(buf.String(), com.Args()...)
	if err != nil {
		return nil, err
	}

	return &firstIterator{rows: rows}, nil
}
func (r *firstRepositoryBase) insert(e *firstEntity) (*firstEntity, error) {
		insert := pqcomp.New(0, 2)
	insert.AddExpr(tableFirstColumnName, "", e.name)

		b := bytes.NewBufferString("INSERT INTO " + r.table)

		if insert.Len() != 0 {
			b.WriteString(" (")
			for insert.Next() {
				if !insert.First() {
					b.WriteString(", ")
				}

				fmt.Fprintf(b, "%s", insert.Key())
			}
			insert.Reset()
			b.WriteString(") VALUES (")
			for insert.Next() {
				if !insert.First() {
					b.WriteString(", ")
				}

				fmt.Fprintf(b, "%s", insert.PlaceHolder())
			}
			b.WriteString(")")
			if len(r.columns) > 0 {
				b.WriteString(" RETURNING ")
				b.WriteString(strings.Join(r.columns, ", "))
			}
		}

		if r.dbg {
			if err := r.log.Log("msg", b.String(), "function", "Insert"); err != nil {
				return nil, err
			}
		}

		err := r.db.QueryRow(b.String(), insert.Args()...).Scan(
	&e.id,
&e.name,
)
		if err != nil {
			return nil, err
		}

		return e, nil
	}
func (r *firstRepositoryBase) upsert(e *firstEntity, p *firstPatch, inf ...string) (*firstEntity, error) {
		insert := pqcomp.New(0, 2)
		update := insert.Compose(2)
	insert.AddExpr(tableFirstColumnName, "", e.name)
if len(inf) > 0 {
update.AddExpr(tableFirstColumnName, "=", p.name)
}

		b := bytes.NewBufferString("INSERT INTO " + r.table)

		if insert.Len() > 0 {
			b.WriteString(" (")
			for insert.Next() {
				if !insert.First() {
					b.WriteString(", ")
				}

				fmt.Fprintf(b, "%s", insert.Key())
			}
			insert.Reset()
			b.WriteString(") VALUES (")
			for insert.Next() {
				if !insert.First() {
					b.WriteString(", ")
				}

				fmt.Fprintf(b, "%s", insert.PlaceHolder())
			}
			b.WriteString(")")
		}
		b.WriteString(" ON CONFLICT ")
		if len(inf) > 0 && update.Len() > 0 {
			b.WriteString(" (")
			for j, i := range inf {
				if j != 0 {
					b.WriteString(", ")
				}
				b.WriteString(i)
			}
			b.WriteString(") ")
			b.WriteString(" DO UPDATE SET ")
			for update.Next() {
				if !update.First() {
					b.WriteString(", ")
				}

				b.WriteString(update.Key())
				b.WriteString(" ")
				b.WriteString(update.Oper())
				b.WriteString(" ")
				b.WriteString(update.PlaceHolder())
			}
		} else {
			b.WriteString(" DO NOTHING ")
		}
		if insert.Len() > 0 {
			if len(r.columns) > 0 {
				b.WriteString(" RETURNING ")
				b.WriteString(strings.Join(r.columns, ", "))
			}
		}

		if r.dbg {
			if err := r.log.Log("msg", b.String(), "function", "Upsert"); err != nil {
				return nil, err
			}
		}

		err := r.db.QueryRow(b.String(), insert.Args()...).Scan(
	&e.id,
&e.name,
)
		if err != nil {
			return nil, err
		}

		return e, nil
	}
`,
		},
	}

	for hint, c := range cases {
		b, err := c.generator.Generate(c.schema)
		if err != nil {
			t.Errorf("%s: unexpected error: %s", hint, err.Error())
			continue
		}
		assertGoCode(t, c.expected, string(b), hint)
	}
}

func assertGoCode(t *testing.T, s1, s2, msg string, com ...interface{}) {
	s1 = fmt.Sprintf("%s", s1)
	s2 = fmt.Sprintf("%s", s2)
	tmp1 := strings.Split(s1, "\n")
	tmp2 := strings.Split(s2, "\n")
	if s1 != s2 {
		b := bytes.NewBuffer(nil)
		for _, diff := range difflib.Diff(tmp1, tmp2) {
			p := strings.Replace(diff.Payload, "\t", "\\t", -1)
			switch diff.Delta {
			case difflib.Common:
				fmt.Fprintf(b, "%s %s\n", diff.Delta.String(), p)
			case difflib.LeftOnly:
				fmt.Fprintf(b, "\033[31m%s %s\033[39m\n", diff.Delta.String(), p)
			case difflib.RightOnly:
				fmt.Fprintf(b, "\033[32m%s %s\033[39m\n", diff.Delta.String(), p)
			}
		}
		t.Errorf(b.String())
	}
}
<|endoftext|>


<commit_before>package models

import (
	"github.com/UniversityRadioYork/myradio-go"
)

// OnDemandModel is the model for the OnDemand controller.
type OnDemandModel struct {
	Model
}

// NewOnDemandModel returns a new OnDemandModel on the MyRadio session s.
func NewOnDemandModel(s *myradio.Session) *OnDemandModel {
	// @TODO: Pass in the config options
	return &OnDemandModel{Model{session: s}}
}

// GetLastMixcloudTimeslots gets the data required for the OnDemand controller from MyRadio.
//
// On success, it returns the previous 3 timeslots, and nil.
// Otherwise, it returns undefined data and the error causing failure.
func (m *OnDemandModel) GetLastMixcloudTimeslots() (timeslots []myradio.Timeslot, err error) {
	timeslots, err = m.session.GetPreviousTimeslots(5)
	if err != nil {
		return
	}

	return
}
<commit_msg>get more timeslots<commit_after>package models

import (
	"github.com/UniversityRadioYork/myradio-go"
)

// OnDemandModel is the model for the OnDemand controller.
type OnDemandModel struct {
	Model
}

// NewOnDemandModel returns a new OnDemandModel on the MyRadio session s.
func NewOnDemandModel(s *myradio.Session) *OnDemandModel {
	// @TODO: Pass in the config options
	return &OnDemandModel{Model{session: s}}
}

// GetLastMixcloudTimeslots gets the data required for the OnDemand controller from MyRadio.
//
// On success, it returns the previous 12 timeslots, and nil.
// Otherwise, it returns undefined data and the error causing failure.
func (m *OnDemandModel) GetLastMixcloudTimeslots() (timeslots []myradio.Timeslot, err error) {
	timeslots, err = m.session.GetPreviousTimeslots(12)
	if err != nil {
		return
	}

	return
}
<|endoftext|>


<commit_before>package user

import (
	"crypto/sha1"
	"errors"
	"fmt"
	"strings"
	"time"

	"github.com/13pinj/todoapp/Godeps/_workspace/src/github.com/gin-gonic/gin"
	"github.com/13pinj/todoapp/Godeps/_workspace/src/github.com/jinzhu/gorm"
	"github.com/13pinj/todoapp/models"
	"github.com/13pinj/todoapp/models/session"
	"github.com/13pinj/todoapp/models/todolist"
)

// Роли пользователя.
// Возможные значения User.Role.
const (
	DefaultRole = ""
	AdminRole   = "admin"
)

// User - структура модели пользователя
type User struct {
	gorm.Model
	Name      string
	PwdHash   string
	Role      string
	Lists     []*todolist.TodoList `gorm:"-"`
	VisitedAt time.Time
}

func hashPwd(pwd string) string {
	return fmt.Sprintf("%x", sha1.Sum([]byte(pwd)))
}

// Register добавляет нового пользователя в базу, и возвращает его структуру,
// если введенные поля корректны. В противном случае Register возвращает ошибку.
func Register(name string, password string) (u *User, errs []error) {
	if strings.Contains(name, " ") {
		errs = append(errs, errors.New("Пробел в имени запрещен"))
	}
	if len([]rune(name)) < 4 {
		errs = append(errs, errors.New("Имя слишком короткое (минимум 4 символа)"))
	}
	if len([]rune(password)) < 6 {
		errs = append(errs, errors.New("Пароль слишком короткий (минимум 6 символов)"))
	}
	if _, ok := Find(name); ok {
		errs = append(errs, errors.New("Имя кем-то занято"))
	}
	if errs != nil {
		return
	}
	u = &User{
		Name:    name,
		PwdHash: hashPwd(password),
	}
	err := models.DB.Save(u).Error
	if err != nil {
		errs = append(errs, err)
	}
	return
}

// Login выполняет авторизацию пользователей.
// Если введенные имя и пароль действительны, Login запишет факт авторизации
// в сессию пользователя и вернет первым значением структуру пользователя,
// а вторым true. В противном случае - nil и false.
// Login перезапишет старые данные об авторизации, если таковые имеются.
func Login(c *gin.Context, name string, password string) (*User, bool) {
	user := &User{}
	err := models.DB.Where("name = ?", name).First(user).Error
	if err != nil {
		return nil, false
	}
	if hashPwd(password) != user.PwdHash {
		return nil, false
	}
	st := session.FromContext(c)
	st.SetInt("user_id", int(user.ID))
	return user, true
}

// Logout стирает данные об авторизации из сессии пользователя.
func Logout(c *gin.Context) {
	st := session.FromContext(c)
	st.SetInt("user_id", 0)
}

// FromContext получает данные об авторизации из сессии пользователя.
// Если пользователь авторизован, FromContext вернет структуру и true.
// Иначе nil и false.
func FromContext(c *gin.Context) (*User, bool) {
	st := session.FromContext(c)
	if st.GetInt("user_id") == 0 {
		return nil, false
	}
	userid := uint(st.GetInt("user_id"))
	user := &User{}
	err := models.DB.Find(user, userid).Error
	if err != nil {
		return nil, false
	}
	user.MarkVisit()
	return user, true
}

// Find находит пользователя в базе по указанному имени.
// Второе возвращаемое значение будет равно false в случае безуспешного поиска.
func Find(name string) (*User, bool) {
	user := &User{}
	err := models.DB.Where("name = ?", name).First(user).Error
	if err != nil {
		return nil, false
	}
	return user, true
}

// AutoLogin запишет факт авторизации в сессию пользователя.
// Он перезапишет старые данные об авторизации, если таковые имеются.
func (u *User) AutoLogin(c *gin.Context) {
	st := session.FromContext(c)
	st.SetInt("user_id", int(u.ID))
}

// MarkVisit обновляет поле VisitedAt и сохраняет в базу.
func (u *User) MarkVisit() {
	u.VisitedAt = time.Now()
	models.DB.Save(u)
}

func (u *User) Online() bool {
	return time.Now().Sub(u.VisitedAt) < 15*time.Minute
}

// LoadLists загружает из базы списки дел пользователя в поле Lists
func (u *User) LoadLists() {
	if u.Lists != nil {
		return
	}
	u.Lists = todolist.FindByUser(u.ID)
}

// SetRole задает пользователю новую роль и сохраняют в базу.
func (u *User) SetRole(r string) {
	u.Role = r
	models.DB.Save(u)
}

// Admin возвращает true, если пользователь относится к администрации.
func (u *User) Admin() bool {
	return u.Role == AdminRole
}

func (u *User) AdminPath() string {
	return fmt.Sprintf("/admin/u/%v", u.Name)
}

func (u *User) CountLists() (n int) {
	models.DB.Model(&todolist.TodoList{}).Where("user_id = ?", u.ID).Count(&n)
	return
}

// Destroy стирает данные о пользователе из базы данных.
func (u *User) Destroy() {
	models.DB.Delete(u)
}

// Count возвращает общее количество всех существующих пользователей.
func Count() (n int) {
	models.DB.Model(&User{}).Count(&n)
	return
}

// Pages возвращает количество страниц, на которые мог бы поместиться
// список всех существующих пользователей по n элементов на страницу.
func Pages(n int) int {
	return (Count()-1)/n + 1
}

// SortMode определяет режим сортировки выборки пользователей.
type SortMode string

// Возможные варианты сортировки выборок пользователей.
const (
	ByID            SortMode = "id"
	ByName          SortMode = "name"
	ByCreatedAt     SortMode = "created_at"
	ByVisitedAt     SortMode = "visited_at, created_at"
	ByIDDesc        SortMode = "id desc"
	ByNameDesc      SortMode = "name desc"
	ByCreatedAtDesc SortMode = "created_at desc"
	ByVisitedAtDesc SortMode = "visited_at desc, created_at desc"
)

// FindPage возвращает список пользователей на i-й странице, если бы они
// размещались по n штук на страницу, отсортированные по sortBy.
// Отсчет страниц ведется с единицы.
func FindPage(i, n int, sortBy SortMode) (us []*User) {
	models.DB.Limit(n).Model(&User{}).Offset(n * (i - 1)).Order(string(sortBy)).Find(&us)
	return
}

var initUser = &User{
	Name:    "root",
	PwdHash: hashPwd("12345678"),
	Role:    AdminRole,
}

func init() {
	initializeUsers()
}

func initializeUsers() {
	dummy := &User{}
	if !models.DB.HasTable(dummy) {
		err := models.DB.CreateTable(dummy).Error
		if err != nil {
			panic(err)
		}
		err = models.DB.Create(initUser).Error
		if err != nil {
			panic(err)
		}
	}
	models.DB.AutoMigrate(dummy)
}
<commit_msg>Set VisitedAt on user registration<commit_after>package user

import (
	"crypto/sha1"
	"errors"
	"fmt"
	"strings"
	"time"

	"github.com/13pinj/todoapp/Godeps/_workspace/src/github.com/gin-gonic/gin"
	"github.com/13pinj/todoapp/Godeps/_workspace/src/github.com/jinzhu/gorm"
	"github.com/13pinj/todoapp/models"
	"github.com/13pinj/todoapp/models/session"
	"github.com/13pinj/todoapp/models/todolist"
)

// Роли пользователя.
// Возможные значения User.Role.
const (
	DefaultRole = ""
	AdminRole   = "admin"
)

// User - структура модели пользователя
type User struct {
	gorm.Model
	Name      string
	PwdHash   string
	Role      string
	Lists     []*todolist.TodoList `gorm:"-"`
	VisitedAt time.Time
}

func hashPwd(pwd string) string {
	return fmt.Sprintf("%x", sha1.Sum([]byte(pwd)))
}

// Register добавляет нового пользователя в базу, и возвращает его структуру,
// если введенные поля корректны. В противном случае Register возвращает ошибку.
func Register(name string, password string) (u *User, errs []error) {
	if strings.Contains(name, " ") {
		errs = append(errs, errors.New("Пробел в имени запрещен"))
	}
	if len([]rune(name)) < 4 {
		errs = append(errs, errors.New("Имя слишком короткое (минимум 4 символа)"))
	}
	if len([]rune(password)) < 6 {
		errs = append(errs, errors.New("Пароль слишком короткий (минимум 6 символов)"))
	}
	if _, ok := Find(name); ok {
		errs = append(errs, errors.New("Имя кем-то занято"))
	}
	if errs != nil {
		return
	}
	u = &User{
		Name:      name,
		PwdHash:   hashPwd(password),
		VisitedAt: time.Now(),
	}
	err := models.DB.Save(u).Error
	if err != nil {
		errs = append(errs, err)
	}
	return
}

// Login выполняет авторизацию пользователей.
// Если введенные имя и пароль действительны, Login запишет факт авторизации
// в сессию пользователя и вернет первым значением структуру пользователя,
// а вторым true. В противном случае - nil и false.
// Login перезапишет старые данные об авторизации, если таковые имеются.
func Login(c *gin.Context, name string, password string) (*User, bool) {
	user := &User{}
	err := models.DB.Where("name = ?", name).First(user).Error
	if err != nil {
		return nil, false
	}
	if hashPwd(password) != user.PwdHash {
		return nil, false
	}
	st := session.FromContext(c)
	st.SetInt("user_id", int(user.ID))
	return user, true
}

// Logout стирает данные об авторизации из сессии пользователя.
func Logout(c *gin.Context) {
	st := session.FromContext(c)
	st.SetInt("user_id", 0)
}

// FromContext получает данные об авторизации из сессии пользователя.
// Если пользователь авторизован, FromContext вернет структуру и true.
// Иначе nil и false.
func FromContext(c *gin.Context) (*User, bool) {
	st := session.FromContext(c)
	if st.GetInt("user_id") == 0 {
		return nil, false
	}
	userid := uint(st.GetInt("user_id"))
	user := &User{}
	err := models.DB.Find(user, userid).Error
	if err != nil {
		return nil, false
	}
	user.MarkVisit()
	return user, true
}

// Find находит пользователя в базе по указанному имени.
// Второе возвращаемое значение будет равно false в случае безуспешного поиска.
func Find(name string) (*User, bool) {
	user := &User{}
	err := models.DB.Where("name = ?", name).First(user).Error
	if err != nil {
		return nil, false
	}
	return user, true
}

// AutoLogin запишет факт авторизации в сессию пользователя.
// Он перезапишет старые данные об авторизации, если таковые имеются.
func (u *User) AutoLogin(c *gin.Context) {
	st := session.FromContext(c)
	st.SetInt("user_id", int(u.ID))
}

// MarkVisit обновляет поле VisitedAt и сохраняет в базу.
func (u *User) MarkVisit() {
	u.VisitedAt = time.Now()
	models.DB.Save(u)
}

func (u *User) Online() bool {
	return time.Now().Sub(u.VisitedAt) < 15*time.Minute
}

// LoadLists загружает из базы списки дел пользователя в поле Lists
func (u *User) LoadLists() {
	if u.Lists != nil {
		return
	}
	u.Lists = todolist.FindByUser(u.ID)
}

// SetRole задает пользователю новую роль и сохраняют в базу.
func (u *User) SetRole(r string) {
	u.Role = r
	models.DB.Save(u)
}

// Admin возвращает true, если пользователь относится к администрации.
func (u *User) Admin() bool {
	return u.Role == AdminRole
}

func (u *User) AdminPath() string {
	return fmt.Sprintf("/admin/u/%v", u.Name)
}

func (u *User) CountLists() (n int) {
	models.DB.Model(&todolist.TodoList{}).Where("user_id = ?", u.ID).Count(&n)
	return
}

// Destroy стирает данные о пользователе из базы данных.
func (u *User) Destroy() {
	models.DB.Delete(u)
}

// Count возвращает общее количество всех существующих пользователей.
func Count() (n int) {
	models.DB.Model(&User{}).Count(&n)
	return
}

// Pages возвращает количество страниц, на которые мог бы поместиться
// список всех существующих пользователей по n элементов на страницу.
func Pages(n int) int {
	return (Count()-1)/n + 1
}

// SortMode определяет режим сортировки выборки пользователей.
type SortMode string

// Возможные варианты сортировки выборок пользователей.
const (
	ByID            SortMode = "id"
	ByName          SortMode = "name"
	ByCreatedAt     SortMode = "created_at"
	ByVisitedAt     SortMode = "visited_at, created_at"
	ByIDDesc        SortMode = "id desc"
	ByNameDesc      SortMode = "name desc"
	ByCreatedAtDesc SortMode = "created_at desc"
	ByVisitedAtDesc SortMode = "visited_at desc, created_at desc"
)

// FindPage возвращает список пользователей на i-й странице, если бы они
// размещались по n штук на страницу, отсортированные по sortBy.
// Отсчет страниц ведется с единицы.
func FindPage(i, n int, sortBy SortMode) (us []*User) {
	models.DB.Limit(n).Model(&User{}).Offset(n * (i - 1)).Order(string(sortBy)).Find(&us)
	return
}

var initUser = &User{
	Name:    "root",
	PwdHash: hashPwd("12345678"),
	Role:    AdminRole,
}

func init() {
	initializeUsers()
}

func initializeUsers() {
	dummy := &User{}
	if !models.DB.HasTable(dummy) {
		err := models.DB.CreateTable(dummy).Error
		if err != nil {
			panic(err)
		}
		err = models.DB.Create(initUser).Error
		if err != nil {
			panic(err)
		}
	}
	models.DB.AutoMigrate(dummy)
}
<|endoftext|>


<commit_before>// Copyright 2014-2015 Amazon.com, Inc. or its affiliates. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License"). You may
// not use this file except in compliance with the License. A copy of the
// License is located at
//
//	http://aws.amazon.com/apache2.0/
//
// or in the "license" file accompanying this file. This file 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 wsclient wraps the generated aws-sdk-go client to provide marshalling
// and unmarshalling of data over a websocket connection in the format expected
// by backend. It allows for bidirectional communication and acts as both a
// client-and-server in terms of requests, but only as a client in terms of
// connecting.
package wsclient

import (
	"crypto/tls"
	"encoding/json"
	"errors"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"reflect"
	"strings"
	"time"

	"github.com/aws/amazon-ecs-agent/agent/logger"
	"github.com/aws/amazon-ecs-agent/agent/utils"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/internal/protocol/json/jsonutil"
	"github.com/gorilla/websocket"
)

var log = logger.ForModule("ws client")

const (
	// ServiceName defines the service name for the agent. This is used to sign messages
	// that are sent to the backend.
	ServiceName = "ecs"

	// wsConnectTimeout specifies the default connection timeout to the backend.
	wsConnectTimeout = 3 * time.Second

	// readBufSize is the size of the read buffer for the ws connection.
	readBufSize = 4096

	// writeBufSize is the size of the write buffer for the ws connection.
	writeBufSize = 32768
)

// ReceivedMessage is the intermediate message used to unmarshal a
// message from backend
type ReceivedMessage struct {
	Type    string          `json:"type"`
	Message json.RawMessage `json:"message"`
}

// RequestMessage is the intermediate message marshalled to send to backend.
type RequestMessage struct {
	Type    string          `json:"type"`
	Message json.RawMessage `json:"message"`
}

// WebsocketConn specifies the subset of gorilla/websocket's
// connection's methods that this client uses.
type WebsocketConn interface {
	WriteMessage(messageType int, data []byte) error
	ReadMessage() (messageType int, data []byte, err error)
	Close() error
}

// RequestHandler would be func(*ecsacs.T for T in ecsacs.*) to be more proper, but it needs
// to be interface{} to properly capture that
type RequestHandler interface{}

// ClientServer is a combined client and server for the backend websocket connection
type ClientServer interface {
	AddRequestHandler(RequestHandler)
	// SetAnyRequestHandler takes a function with the signature 'func(i
	// interface{})' and calls it with every message the server passes down.
	// Only a single 'AnyRequestHandler' will be active at a given time for a
	// ClientServer
	SetAnyRequestHandler(RequestHandler)
	MakeRequest(input interface{}) error
	Connect() error
	Serve() error
	io.Closer
}

//go:generate mockgen.sh github.com/aws/amazon-ecs-agent/agent/wsclient ClientServer mock/$GOFILE

// ClientServerImpl wraps commonly used methods defined in ClientServer interface.
type ClientServerImpl struct {
	AcceptInvalidCert  bool
	Conn               WebsocketConn
	CredentialProvider *credentials.Credentials
	Region             string
	// RequestHandlers is a map from message types to handler functions of the
	// form:
	//     "FooMessage": func(message *ecsacs.FooMessage)
	RequestHandlers map[string]RequestHandler
	// AnyRequestHandler is a request handler that, if set, is called on every
	// message with said message. It will be called before a RequestHandler is
	// called. It must take a single interface{} argument.
	AnyRequestHandler RequestHandler
	// URL is the full url to the backend, including path, querystring, and so on.
	URL string
	ClientServer
	ServiceError
	TypeDecoder
}

// Connect opens a connection to the backend and upgrades it to a websocket. Calls to
// 'MakeRequest' can be made after calling this, but responss will not be
// receivable until 'Serve' is also called.
func (cs *ClientServerImpl) Connect() error {
	parsedURL, err := url.Parse(cs.URL)
	if err != nil {
		return err
	}

	// NewRequest never returns an error if the url parses and we just verified
	// it did above
	request, _ := http.NewRequest("GET", cs.URL, nil)
	// Sign the request; we'll send its headers via the websocket client which includes the signature
	utils.SignHTTPRequest(request, cs.Region, ServiceName, cs.CredentialProvider, nil)

	// url.Host might not have the port, but tls.Dial needs it
	dialHost := parsedURL.Host
	if !strings.Contains(dialHost, ":") {
		dialHost += ":443"
	}

	timeoutDialer := &net.Dialer{Timeout: wsConnectTimeout}
	log.Info("Creating poll dialer", "host", parsedURL.Host)
	wsConn, err := tls.DialWithDialer(timeoutDialer, "tcp", dialHost, &tls.Config{InsecureSkipVerify: cs.AcceptInvalidCert})
	if err != nil {
		return err
	}

	websocketConn, httpResponse, err := websocket.NewClient(wsConn, parsedURL, request.Header, readBufSize, writeBufSize)
	if httpResponse != nil {
		defer httpResponse.Body.Close()
	}
	if err != nil {
		defer wsConn.Close()
		var resp []byte
		if httpResponse != nil {
			var readErr error
			resp, readErr = ioutil.ReadAll(httpResponse.Body)
			if readErr != nil {
				return errors.New("Unable to read websocket connection: " + readErr.Error() + ", " + err.Error())
			}
			// If there's a response, we can try to unmarshal it into one of the
			// modeled error types
			possibleError, _, decodeErr := DecodeData(resp, cs.TypeDecoder)
			if decodeErr == nil {
				return cs.NewError(possibleError)
			}
		}
		log.Warn("Error creating a websocket client", "err", err)
		return errors.New(string(resp) + ", " + err.Error())
	}
	cs.Conn = websocketConn
	return nil
}

// AddRequestHandler adds a request handler to this client.
// A request handler *must* be a function taking a single argument, and that
// argument *must* be a pointer to a recognized 'ecsacs' struct.
// E.g. if you desired to handle messages from acs of type 'FooMessage', you
// would pass the following handler in:
//     func(message *ecsacs.FooMessage)
// This function will panic if the passed in function does not have one pointer
// argument or the argument is not a recognized type.
// Additionally, the request handler will block processing of further messages
// on this connection so it's important that it return quickly.
func (cs *ClientServerImpl) AddRequestHandler(f RequestHandler) {
	firstArg := reflect.TypeOf(f).In(0)
	firstArgTypeStr := firstArg.Elem().Name()
	recognizedTypes := cs.GetRecognizedTypes()
	_, ok := recognizedTypes[firstArgTypeStr]
	if !ok {
		panic("AddRequestHandler called with invalid function; argument type not recognized: " + firstArgTypeStr)
	}
	cs.RequestHandlers[firstArgTypeStr] = f
}

func (cs *ClientServerImpl) SetAnyRequestHandler(f RequestHandler) {
	cs.AnyRequestHandler = f
}

// MakeRequest makes a request using the given input. Note, the input *MUST* be
// a pointer to a valid backend type that this client recognises
func (cs *ClientServerImpl) MakeRequest(input interface{}) error {
	send, err := cs.CreateRequestMessage(input)
	if err != nil {
		return err
	}

	// Over the wire we send something like
	// {"type":"AckRequest","message":{"messageId":"xyz"}}
	return cs.Conn.WriteMessage(websocket.TextMessage, send)
}

// ConsumeMessages reads messages from the websocket connection and handles read
// messages from an active connection.
func (cs *ClientServerImpl) ConsumeMessages() error {
	var err error
	for {
		messageType, message, cerr := cs.Conn.ReadMessage()
		err = cerr
		if err != nil {
			if err != io.EOF {
				if message != nil {
					log.Error("Error getting message from ws backend", "err", err, "message", message)
				} else {
					log.Error("Error getting message from ws backend", "err", err)
				}
			}
			break
		}
		if messageType != websocket.TextMessage {
			log.Error("Unexpected messageType", "type", messageType)
			// maybe not fatal though, we'll try to process it anyways
		}
		log.Debug("Got a message from websocket", "message", string(message[:]))
		cs.handleMessage(message)
	}
	return err
}

// CreateRequestMessage creates the request json message using the given input.
// Note, the input *MUST* be a pointer to a valid backend type that this
// client recognises.
func (cs *ClientServerImpl) CreateRequestMessage(input interface{}) ([]byte, error) {
	msg := &RequestMessage{}

	recognizedTypes := cs.GetRecognizedTypes()
	for typeStr, typeVal := range recognizedTypes {
		if reflect.TypeOf(input) == reflect.PtrTo(typeVal) {
			msg.Type = typeStr
			break
		}
	}
	if msg.Type == "" {
		return nil, &UnrecognizedWSRequestType{reflect.TypeOf(input).String()}
	}
	messageData, err := jsonutil.BuildJSON(input)
	if err != nil {
		return nil, &NotMarshallableWSRequest{msg.Type, err}
	}
	msg.Message = json.RawMessage(messageData)

	send, err := json.Marshal(msg)
	if err != nil {
		return nil, &NotMarshallableWSRequest{msg.Type, err}
	}
	return send, nil
}

// handleMessage dispatches a message to the correct 'requestHandler' for its
// type. If no request handler is found, the message is discarded.
func (cs *ClientServerImpl) handleMessage(data []byte) {
	typedMessage, typeStr, err := DecodeData(data, cs.TypeDecoder)
	if err != nil {
		log.Warn("Unable to handle message from backend", "err", err)
		return
	}

	if cs.AnyRequestHandler != nil {
		reflect.ValueOf(cs.AnyRequestHandler).Call([]reflect.Value{reflect.ValueOf(typedMessage)})
	}

	if handler, ok := cs.RequestHandlers[typeStr]; ok {
		reflect.ValueOf(handler).Call([]reflect.Value{reflect.ValueOf(typedMessage)})
	} else {
		log.Info("No handler for message type", "type", typeStr)
	}
}
<commit_msg>Added support for proxied connections to wsclient<commit_after>// Copyright 2014-2015 Amazon.com, Inc. or its affiliates. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License"). You may
// not use this file except in compliance with the License. A copy of the
// License is located at
//
//	http://aws.amazon.com/apache2.0/
//
// or in the "license" file accompanying this file. This file 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 wsclient wraps the generated aws-sdk-go client to provide marshalling
// and unmarshalling of data over a websocket connection in the format expected
// by backend. It allows for bidirectional communication and acts as both a
// client-and-server in terms of requests, but only as a client in terms of
// connecting.
package wsclient

import (
	"bufio"
	"crypto/tls"
	"encoding/base64"
	"encoding/json"
	"errors"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"reflect"
	"strings"
	"time"

	"github.com/aws/amazon-ecs-agent/agent/logger"
	"github.com/aws/amazon-ecs-agent/agent/utils"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/internal/protocol/json/jsonutil"
	"github.com/gorilla/websocket"
)

var log = logger.ForModule("ws client")

const (
	// ServiceName defines the service name for the agent. This is used to sign messages
	// that are sent to the backend.
	ServiceName = "ecs"

	// wsConnectTimeout specifies the default connection timeout to the backend.
	wsConnectTimeout = 3 * time.Second

	// readBufSize is the size of the read buffer for the ws connection.
	readBufSize = 4096

	// writeBufSize is the size of the write buffer for the ws connection.
	writeBufSize = 32768
)

// ReceivedMessage is the intermediate message used to unmarshal a
// message from backend
type ReceivedMessage struct {
	Type    string          `json:"type"`
	Message json.RawMessage `json:"message"`
}

// RequestMessage is the intermediate message marshalled to send to backend.
type RequestMessage struct {
	Type    string          `json:"type"`
	Message json.RawMessage `json:"message"`
}

// WebsocketConn specifies the subset of gorilla/websocket's
// connection's methods that this client uses.
type WebsocketConn interface {
	WriteMessage(messageType int, data []byte) error
	ReadMessage() (messageType int, data []byte, err error)
	Close() error
}

// RequestHandler would be func(*ecsacs.T for T in ecsacs.*) to be more proper, but it needs
// to be interface{} to properly capture that
type RequestHandler interface{}

// ClientServer is a combined client and server for the backend websocket connection
type ClientServer interface {
	AddRequestHandler(RequestHandler)
	// SetAnyRequestHandler takes a function with the signature 'func(i
	// interface{})' and calls it with every message the server passes down.
	// Only a single 'AnyRequestHandler' will be active at a given time for a
	// ClientServer
	SetAnyRequestHandler(RequestHandler)
	MakeRequest(input interface{}) error
	Connect() error
	Serve() error
	io.Closer
}

//go:generate mockgen.sh github.com/aws/amazon-ecs-agent/agent/wsclient ClientServer mock/$GOFILE

// ClientServerImpl wraps commonly used methods defined in ClientServer interface.
type ClientServerImpl struct {
	AcceptInvalidCert  bool
	Conn               WebsocketConn
	CredentialProvider *credentials.Credentials
	Region             string
	// RequestHandlers is a map from message types to handler functions of the
	// form:
	//     "FooMessage": func(message *ecsacs.FooMessage)
	RequestHandlers map[string]RequestHandler
	// AnyRequestHandler is a request handler that, if set, is called on every
	// message with said message. It will be called before a RequestHandler is
	// called. It must take a single interface{} argument.
	AnyRequestHandler RequestHandler
	// URL is the full url to the backend, including path, querystring, and so on.
	URL string
	ClientServer
	ServiceError
	TypeDecoder
}

// Connect opens a connection to the backend and upgrades it to a websocket. Calls to
// 'MakeRequest' can be made after calling this, but responss will not be
// receivable until 'Serve' is also called.
func (cs *ClientServerImpl) Connect() error {
	parsedURL, err := url.Parse(cs.URL)
	if err != nil {
		return err
	}

	// NewRequest never returns an error if the url parses and we just verified
	// it did above
	request, _ := http.NewRequest("GET", cs.URL, nil)
	// Sign the request; we'll send its headers via the websocket client which includes the signature
	utils.SignHTTPRequest(request, cs.Region, ServiceName, cs.CredentialProvider, nil)

	wsConn, err := cs.websocketConn(parsedURL, request)
	if err != nil {
		return err
	}

	websocketConn, httpResponse, err := websocket.NewClient(wsConn, parsedURL, request.Header, readBufSize, writeBufSize)
	if httpResponse != nil {
		defer httpResponse.Body.Close()
	}
	if err != nil {
		defer wsConn.Close()
		var resp []byte
		if httpResponse != nil {
			var readErr error
			resp, readErr = ioutil.ReadAll(httpResponse.Body)
			if readErr != nil {
				return errors.New("Unable to read websocket connection: " + readErr.Error() + ", " + err.Error())
			}
			// If there's a response, we can try to unmarshal it into one of the
			// modeled error types
			possibleError, _, decodeErr := DecodeData(resp, cs.TypeDecoder)
			if decodeErr == nil {
				return cs.NewError(possibleError)
			}
		}
		log.Warn("Error creating a websocket client", "err", err)
		return errors.New(string(resp) + ", " + err.Error())
	}
	cs.Conn = websocketConn
	return nil
}

// AddRequestHandler adds a request handler to this client.
// A request handler *must* be a function taking a single argument, and that
// argument *must* be a pointer to a recognized 'ecsacs' struct.
// E.g. if you desired to handle messages from acs of type 'FooMessage', you
// would pass the following handler in:
//     func(message *ecsacs.FooMessage)
// This function will panic if the passed in function does not have one pointer
// argument or the argument is not a recognized type.
// Additionally, the request handler will block processing of further messages
// on this connection so it's important that it return quickly.
func (cs *ClientServerImpl) AddRequestHandler(f RequestHandler) {
	firstArg := reflect.TypeOf(f).In(0)
	firstArgTypeStr := firstArg.Elem().Name()
	recognizedTypes := cs.GetRecognizedTypes()
	_, ok := recognizedTypes[firstArgTypeStr]
	if !ok {
		panic("AddRequestHandler called with invalid function; argument type not recognized: " + firstArgTypeStr)
	}
	cs.RequestHandlers[firstArgTypeStr] = f
}

func (cs *ClientServerImpl) SetAnyRequestHandler(f RequestHandler) {
	cs.AnyRequestHandler = f
}

// MakeRequest makes a request using the given input. Note, the input *MUST* be
// a pointer to a valid backend type that this client recognises
func (cs *ClientServerImpl) MakeRequest(input interface{}) error {
	send, err := cs.CreateRequestMessage(input)
	if err != nil {
		return err
	}

	// Over the wire we send something like
	// {"type":"AckRequest","message":{"messageId":"xyz"}}
	return cs.Conn.WriteMessage(websocket.TextMessage, send)
}

// ConsumeMessages reads messages from the websocket connection and handles read
// messages from an active connection.
func (cs *ClientServerImpl) ConsumeMessages() error {
	var err error
	for {
		messageType, message, cerr := cs.Conn.ReadMessage()
		err = cerr
		if err != nil {
			if err != io.EOF {
				if message != nil {
					log.Error("Error getting message from ws backend", "err", err, "message", message)
				} else {
					log.Error("Error getting message from ws backend", "err", err)
				}
			}
			break
		}
		if messageType != websocket.TextMessage {
			log.Error("Unexpected messageType", "type", messageType)
			// maybe not fatal though, we'll try to process it anyways
		}
		log.Debug("Got a message from websocket", "message", string(message[:]))
		cs.handleMessage(message)
	}
	return err
}

// CreateRequestMessage creates the request json message using the given input.
// Note, the input *MUST* be a pointer to a valid backend type that this
// client recognises.
func (cs *ClientServerImpl) CreateRequestMessage(input interface{}) ([]byte, error) {
	msg := &RequestMessage{}

	recognizedTypes := cs.GetRecognizedTypes()
	for typeStr, typeVal := range recognizedTypes {
		if reflect.TypeOf(input) == reflect.PtrTo(typeVal) {
			msg.Type = typeStr
			break
		}
	}
	if msg.Type == "" {
		return nil, &UnrecognizedWSRequestType{reflect.TypeOf(input).String()}
	}
	messageData, err := jsonutil.BuildJSON(input)
	if err != nil {
		return nil, &NotMarshallableWSRequest{msg.Type, err}
	}
	msg.Message = json.RawMessage(messageData)

	send, err := json.Marshal(msg)
	if err != nil {
		return nil, &NotMarshallableWSRequest{msg.Type, err}
	}
	return send, nil
}

// handleMessage dispatches a message to the correct 'requestHandler' for its
// type. If no request handler is found, the message is discarded.
func (cs *ClientServerImpl) handleMessage(data []byte) {
	typedMessage, typeStr, err := DecodeData(data, cs.TypeDecoder)
	if err != nil {
		log.Warn("Unable to handle message from backend", "err", err)
		return
	}

	if cs.AnyRequestHandler != nil {
		reflect.ValueOf(cs.AnyRequestHandler).Call([]reflect.Value{reflect.ValueOf(typedMessage)})
	}

	if handler, ok := cs.RequestHandlers[typeStr]; ok {
		reflect.ValueOf(handler).Call([]reflect.Value{reflect.ValueOf(typedMessage)})
	} else {
		log.Info("No handler for message type", "type", typeStr)
	}
}

// websocketConn establishes a connection to the given URL, respecting any proxy configuration in the environment.
// A standard proxying setup involves setting the following environment variables
// (may be listed in /etc/ecs/ecs.config if using ecs-init):
// HTTP_PROXY=http://<your-proxy>/ # HTTPS_PROXY may be set instead or additionally
// NO_PROXY=169.254.169.254,/var/run/docker.sock # Directly connect to metadata service and docker socket
func (cs *ClientServerImpl) websocketConn(parsedURL *url.URL, request *http.Request) (net.Conn, error) {
	proxyURL, err := http.ProxyFromEnvironment(request)
	if err != nil {
		return nil, err
	}

	// url.Host might not have the port, but tls.Dial needs it
	targetHost := parsedURL.Host
	if !strings.Contains(targetHost, ":") {
		targetHost += ":443"
	}
	targetHostname, _ , err := net.SplitHostPort(targetHost)
	if err != nil {
		return nil, err
	}

	tlsConfig := tls.Config{ServerName: targetHostname, InsecureSkipVerify: cs.AcceptInvalidCert}
	timeoutDialer := &net.Dialer{Timeout: wsConnectTimeout}

	if proxyURL == nil {
		// directly connect
		log.Info("Creating poll dialer", "host", parsedURL.Host)
		return tls.DialWithDialer(timeoutDialer, "tcp", targetHost, &tlsConfig)
	}

	// connect via proxy
	log.Info("Creating poll dialer", "proxy", proxyURL.Host, "host", parsedURL.Host)
	plainConn, err := timeoutDialer.Dial("tcp", proxyURL.Host)
	if err != nil {
		return nil, err
	}

	// TLS over an HTTP proxy via CONNECT taken from: https://golang.org/src/net/http/transport.go
	connectReq := &http.Request{
		Method: "CONNECT",
		URL: &url.URL{Opaque: targetHost},
		Host: targetHost,
		Header: make(http.Header),
	}

	if proxyUser := proxyURL.User; proxyUser != nil {
		username := proxyUser.Username()
		password, _ := proxyUser.Password()
		auth := base64.StdEncoding.EncodeToString([]byte(username + ":" + password))
		connectReq.Header.Set("Proxy-Authorization", "Basic " + auth)
	}

	connectReq.Write(plainConn)

	// Read response.
	// Okay to use and discard buffered reader here, because
	// TLS server will not speak until spoken to.
	br := bufio.NewReader(plainConn)
	resp, err := http.ReadResponse(br, connectReq)
	if err != nil {
		plainConn.Close()
		return nil, err
	}
	if resp.StatusCode != 200 {
		plainConn.Close()
		return nil, errors.New(resp.Status)
	}

	tlsConn := tls.Client(plainConn, &tlsConfig)

	return tlsConn, nil
}
<|endoftext|>


<commit_before>package peer

import (
	"math"
	"net"
	"time"

	"github.com/cenkalti/rain/internal/bitfield"
	"github.com/cenkalti/rain/internal/fast"
	"github.com/cenkalti/rain/internal/logger"
	"github.com/cenkalti/rain/internal/mse"
	"github.com/cenkalti/rain/internal/peerconn"
	"github.com/cenkalti/rain/internal/peerconn/peerreader"
	"github.com/cenkalti/rain/internal/peerconn/peerwriter"
	"github.com/cenkalti/rain/internal/peerprotocol"
	"github.com/cenkalti/rain/internal/peersource"
	"github.com/cenkalti/rain/internal/pexlist"
	"github.com/cenkalti/rain/internal/piece"
	"github.com/cenkalti/rain/internal/pieceset"
	"github.com/cenkalti/rain/internal/stringutil"
	"github.com/juju/ratelimit"
	"github.com/rcrowley/go-metrics"
)

type Peer struct {
	*peerconn.Conn

	ConnectedAt time.Time

	Source peersource.Source

	Bitfield            *bitfield.Bitfield
	ReceivedAllowedFast pieceset.PieceSet
	SentAllowedFast     pieceset.PieceSet

	ID                [20]byte
	ExtensionsEnabled bool
	FastEnabled       bool
	DHTEnabled        bool
	EncryptionCipher  mse.CryptoMethod

	ClientInterested bool
	ClientChoking    bool
	PeerInterested   bool
	PeerChoking      bool

	OptimisticUnchoked bool

	// Snubbed means peer is sending pieces too slow.
	Snubbed bool

	Downloading bool

	downloadSpeed metrics.Meter
	uploadSpeed   metrics.Meter

	// Messages received while we don't have info yet are saved here.
	Messages []interface{}

	ExtensionHandshake *peerprotocol.ExtensionHandshakeMessage

	PEX *pex

	snubTimeout time.Duration
	snubTimer   *time.Timer

	closeC chan struct{}
	doneC  chan struct{}
}

type Message struct {
	*Peer
	Message interface{}
}

type PieceMessage struct {
	*Peer
	Piece peerreader.Piece
}

func New(conn net.Conn, source peersource.Source, id [20]byte, extensions [8]byte, cipher mse.CryptoMethod, pieceReadTimeout, snubTimeout time.Duration, maxRequestsIn int, br, bw *ratelimit.Bucket) *Peer {
	bf, _ := bitfield.NewBytes(extensions[:], 64)
	fastEnabled := bf.Test(61)
	extensionsEnabled := bf.Test(43)
	dhtEnabled := bf.Test(63)

	t := time.NewTimer(math.MaxInt64)
	t.Stop()
	return &Peer{
		Conn:              peerconn.New(conn, newPeerLogger(source, conn), pieceReadTimeout, maxRequestsIn, fastEnabled, br, bw),
		Source:            source,
		ConnectedAt:       time.Now(),
		ID:                id,
		ClientChoking:     true,
		PeerChoking:       true,
		ExtensionsEnabled: extensionsEnabled,
		FastEnabled:       fastEnabled,
		DHTEnabled:        dhtEnabled,
		EncryptionCipher:  cipher,
		snubTimeout:       snubTimeout,
		snubTimer:         t,
		closeC:            make(chan struct{}),
		doneC:             make(chan struct{}),
		downloadSpeed:     metrics.NewMeter(),
		uploadSpeed:       metrics.NewMeter(),
	}
}

func newPeerLogger(src peersource.Source, conn net.Conn) logger.Logger {
	if src == peersource.Incoming {
		return logger.New("peer <- " + conn.RemoteAddr().String())
	}
	return logger.New("peer -> " + conn.RemoteAddr().String())
}

func (p *Peer) Close() {
	p.snubTimer.Stop()
	if p.PEX != nil {
		p.PEX.close()
	}
	close(p.closeC)
	p.Conn.Close()
	p.downloadSpeed.Stop()
	p.uploadSpeed.Stop()
	<-p.doneC
}

func (p *Peer) Done() chan struct{} {
	return p.doneC
}

func (p *Peer) Run(messages chan Message, pieces chan interface{}, snubbed, disconnect chan *Peer) {
	defer close(p.doneC)
	go p.Conn.Run()

	for {
		select {
		case pm, ok := <-p.Conn.Messages():
			if !ok {
				select {
				case disconnect <- p:
				case <-p.closeC:
				}
				return
			}
			if m, ok := pm.(peerreader.Piece); ok {
				p.downloadSpeed.Mark(int64(len(m.Buffer.Data)))
				select {
				case pieces <- PieceMessage{Peer: p, Piece: m}:
				case <-p.closeC:
					return
				}
			} else {
				if m, ok := pm.(peerwriter.BlockUploaded); ok {
					p.uploadSpeed.Mark(int64(m.Length))
				}
				select {
				case messages <- Message{Peer: p, Message: pm}:
				case <-p.closeC:
					return
				}
			}
		case <-p.snubTimer.C:
			select {
			case snubbed <- p:
			case <-p.closeC:
				return
			}
		case <-p.closeC:
			return
		}
	}
}

func (p *Peer) StartPEX(initialPeers map[*Peer]struct{}, recentlySeen *pexlist.RecentlySeen) {
	if p.PEX == nil {
		p.PEX = newPEX(p.Conn, p.ExtensionHandshake.M[peerprotocol.ExtensionKeyPEX], initialPeers, recentlySeen)
		go p.PEX.run()
	}
}

func (p *Peer) ResetSnubTimer() {
	p.snubTimer.Reset(p.snubTimeout)
}

func (p *Peer) StopSnubTimer() {
	p.snubTimer.Stop()
}

func (p *Peer) DownloadSpeed() int {
	return int(p.downloadSpeed.Rate1())
}

func (p *Peer) UploadSpeed() int {
	return int(p.uploadSpeed.Rate1())
}

func (p *Peer) Choke() {
	p.ClientChoking = true
	p.SendMessage(peerprotocol.ChokeMessage{})
}

func (p *Peer) Unchoke() {
	p.ClientChoking = false
	p.SendMessage(peerprotocol.UnchokeMessage{})
}

func (p *Peer) Choking() bool {
	return p.ClientChoking
}

func (p *Peer) Interested() bool {
	return p.PeerInterested
}

func (p *Peer) Optimistic() bool {
	return p.OptimisticUnchoked
}

func (p *Peer) SetOptimistic(value bool) {
	p.OptimisticUnchoked = value
}

func (p *Peer) MetadataSize() uint32 {
	return uint32(p.ExtensionHandshake.MetadataSize)
}

func (p *Peer) RequestMetadataPiece(index uint32) {
	p.SendMessage(peerprotocol.ExtensionMessage{
		ExtendedMessageID: p.ExtensionHandshake.M[peerprotocol.ExtensionKeyMetadata],
		Payload: peerprotocol.ExtensionMetadataMessage{
			Type:  peerprotocol.ExtensionMetadataMessageTypeRequest,
			Piece: index,
		},
	})
}

func (p *Peer) RequestPiece(index, begin, length uint32) {
	msg := peerprotocol.RequestMessage{Index: index, Begin: begin, Length: length}
	p.SendMessage(msg)
}

func (p *Peer) CancelPiece(index, begin, length uint32) {
	msg := peerprotocol.CancelMessage{RequestMessage: peerprotocol.RequestMessage{Index: index, Begin: begin, Length: length}}
	p.SendMessage(msg)
}

func (p *Peer) EnabledFast() bool {
	return p.FastEnabled
}

func (p *Peer) Client() string {
	if p.ExtensionHandshake == nil {
		return ""
	}
	if p.ExtensionHandshake.V != "" {
		return stringutil.Printable(p.ExtensionHandshake.V)
	}
	return stringutil.Asciify(clientID(string(p.ID[:])))
}

func (p *Peer) GenerateAndSendAllowedFastMessages(k int, numPieces uint32, infoHash [20]byte, pieces []piece.Piece) {
	if k == 0 {
		return
	}
	if p.SentAllowedFast.Len() > 0 {
		return
	}
	a := fast.GenerateFastSet(k, numPieces, infoHash, p.Conn.Addr().IP)
	for _, index := range a {
		p.SentAllowedFast.Add(&pieces[index])
		p.SendMessage(peerprotocol.AllowedFastMessage{HaveMessage: peerprotocol.HaveMessage{Index: index}})
	}
}
<commit_msg>return peer id if no extension handshake received<commit_after>package peer

import (
	"math"
	"net"
	"time"

	"github.com/cenkalti/rain/internal/bitfield"
	"github.com/cenkalti/rain/internal/fast"
	"github.com/cenkalti/rain/internal/logger"
	"github.com/cenkalti/rain/internal/mse"
	"github.com/cenkalti/rain/internal/peerconn"
	"github.com/cenkalti/rain/internal/peerconn/peerreader"
	"github.com/cenkalti/rain/internal/peerconn/peerwriter"
	"github.com/cenkalti/rain/internal/peerprotocol"
	"github.com/cenkalti/rain/internal/peersource"
	"github.com/cenkalti/rain/internal/pexlist"
	"github.com/cenkalti/rain/internal/piece"
	"github.com/cenkalti/rain/internal/pieceset"
	"github.com/cenkalti/rain/internal/stringutil"
	"github.com/juju/ratelimit"
	"github.com/rcrowley/go-metrics"
)

type Peer struct {
	*peerconn.Conn

	ConnectedAt time.Time

	Source peersource.Source

	Bitfield            *bitfield.Bitfield
	ReceivedAllowedFast pieceset.PieceSet
	SentAllowedFast     pieceset.PieceSet

	ID                [20]byte
	ExtensionsEnabled bool
	FastEnabled       bool
	DHTEnabled        bool
	EncryptionCipher  mse.CryptoMethod

	ClientInterested bool
	ClientChoking    bool
	PeerInterested   bool
	PeerChoking      bool

	OptimisticUnchoked bool

	// Snubbed means peer is sending pieces too slow.
	Snubbed bool

	Downloading bool

	downloadSpeed metrics.Meter
	uploadSpeed   metrics.Meter

	// Messages received while we don't have info yet are saved here.
	Messages []interface{}

	ExtensionHandshake *peerprotocol.ExtensionHandshakeMessage

	PEX *pex

	snubTimeout time.Duration
	snubTimer   *time.Timer

	closeC chan struct{}
	doneC  chan struct{}
}

type Message struct {
	*Peer
	Message interface{}
}

type PieceMessage struct {
	*Peer
	Piece peerreader.Piece
}

func New(conn net.Conn, source peersource.Source, id [20]byte, extensions [8]byte, cipher mse.CryptoMethod, pieceReadTimeout, snubTimeout time.Duration, maxRequestsIn int, br, bw *ratelimit.Bucket) *Peer {
	bf, _ := bitfield.NewBytes(extensions[:], 64)
	fastEnabled := bf.Test(61)
	extensionsEnabled := bf.Test(43)
	dhtEnabled := bf.Test(63)

	t := time.NewTimer(math.MaxInt64)
	t.Stop()
	return &Peer{
		Conn:              peerconn.New(conn, newPeerLogger(source, conn), pieceReadTimeout, maxRequestsIn, fastEnabled, br, bw),
		Source:            source,
		ConnectedAt:       time.Now(),
		ID:                id,
		ClientChoking:     true,
		PeerChoking:       true,
		ExtensionsEnabled: extensionsEnabled,
		FastEnabled:       fastEnabled,
		DHTEnabled:        dhtEnabled,
		EncryptionCipher:  cipher,
		snubTimeout:       snubTimeout,
		snubTimer:         t,
		closeC:            make(chan struct{}),
		doneC:             make(chan struct{}),
		downloadSpeed:     metrics.NewMeter(),
		uploadSpeed:       metrics.NewMeter(),
	}
}

func newPeerLogger(src peersource.Source, conn net.Conn) logger.Logger {
	if src == peersource.Incoming {
		return logger.New("peer <- " + conn.RemoteAddr().String())
	}
	return logger.New("peer -> " + conn.RemoteAddr().String())
}

func (p *Peer) Close() {
	p.snubTimer.Stop()
	if p.PEX != nil {
		p.PEX.close()
	}
	close(p.closeC)
	p.Conn.Close()
	p.downloadSpeed.Stop()
	p.uploadSpeed.Stop()
	<-p.doneC
}

func (p *Peer) Done() chan struct{} {
	return p.doneC
}

func (p *Peer) Run(messages chan Message, pieces chan interface{}, snubbed, disconnect chan *Peer) {
	defer close(p.doneC)
	go p.Conn.Run()

	for {
		select {
		case pm, ok := <-p.Conn.Messages():
			if !ok {
				select {
				case disconnect <- p:
				case <-p.closeC:
				}
				return
			}
			if m, ok := pm.(peerreader.Piece); ok {
				p.downloadSpeed.Mark(int64(len(m.Buffer.Data)))
				select {
				case pieces <- PieceMessage{Peer: p, Piece: m}:
				case <-p.closeC:
					return
				}
			} else {
				if m, ok := pm.(peerwriter.BlockUploaded); ok {
					p.uploadSpeed.Mark(int64(m.Length))
				}
				select {
				case messages <- Message{Peer: p, Message: pm}:
				case <-p.closeC:
					return
				}
			}
		case <-p.snubTimer.C:
			select {
			case snubbed <- p:
			case <-p.closeC:
				return
			}
		case <-p.closeC:
			return
		}
	}
}

func (p *Peer) StartPEX(initialPeers map[*Peer]struct{}, recentlySeen *pexlist.RecentlySeen) {
	if p.PEX == nil {
		p.PEX = newPEX(p.Conn, p.ExtensionHandshake.M[peerprotocol.ExtensionKeyPEX], initialPeers, recentlySeen)
		go p.PEX.run()
	}
}

func (p *Peer) ResetSnubTimer() {
	p.snubTimer.Reset(p.snubTimeout)
}

func (p *Peer) StopSnubTimer() {
	p.snubTimer.Stop()
}

func (p *Peer) DownloadSpeed() int {
	return int(p.downloadSpeed.Rate1())
}

func (p *Peer) UploadSpeed() int {
	return int(p.uploadSpeed.Rate1())
}

func (p *Peer) Choke() {
	p.ClientChoking = true
	p.SendMessage(peerprotocol.ChokeMessage{})
}

func (p *Peer) Unchoke() {
	p.ClientChoking = false
	p.SendMessage(peerprotocol.UnchokeMessage{})
}

func (p *Peer) Choking() bool {
	return p.ClientChoking
}

func (p *Peer) Interested() bool {
	return p.PeerInterested
}

func (p *Peer) Optimistic() bool {
	return p.OptimisticUnchoked
}

func (p *Peer) SetOptimistic(value bool) {
	p.OptimisticUnchoked = value
}

func (p *Peer) MetadataSize() uint32 {
	return uint32(p.ExtensionHandshake.MetadataSize)
}

func (p *Peer) RequestMetadataPiece(index uint32) {
	p.SendMessage(peerprotocol.ExtensionMessage{
		ExtendedMessageID: p.ExtensionHandshake.M[peerprotocol.ExtensionKeyMetadata],
		Payload: peerprotocol.ExtensionMetadataMessage{
			Type:  peerprotocol.ExtensionMetadataMessageTypeRequest,
			Piece: index,
		},
	})
}

func (p *Peer) RequestPiece(index, begin, length uint32) {
	msg := peerprotocol.RequestMessage{Index: index, Begin: begin, Length: length}
	p.SendMessage(msg)
}

func (p *Peer) CancelPiece(index, begin, length uint32) {
	msg := peerprotocol.CancelMessage{RequestMessage: peerprotocol.RequestMessage{Index: index, Begin: begin, Length: length}}
	p.SendMessage(msg)
}

func (p *Peer) EnabledFast() bool {
	return p.FastEnabled
}

func (p *Peer) Client() string {
	if p.ExtensionHandshake != nil && p.ExtensionHandshake.V != "" {
		return stringutil.Printable(p.ExtensionHandshake.V)
	}
	return stringutil.Asciify(clientID(string(p.ID[:])))
}

func (p *Peer) GenerateAndSendAllowedFastMessages(k int, numPieces uint32, infoHash [20]byte, pieces []piece.Piece) {
	if k == 0 {
		return
	}
	if p.SentAllowedFast.Len() > 0 {
		return
	}
	a := fast.GenerateFastSet(k, numPieces, infoHash, p.Conn.Addr().IP)
	for _, index := range a {
		p.SentAllowedFast.Add(&pieces[index])
		p.SendMessage(peerprotocol.AllowedFastMessage{HaveMessage: peerprotocol.HaveMessage{Index: index}})
	}
}
<|endoftext|>


<commit_before>package ircfmt

import "testing"

type testcase struct {
	escaped   string
	unescaped string
}

var tests = []testcase{
	{"te$bst", "te\x02st"},
	{"te$c[green]st", "te\x033st"},
	{"te$c[red,green]st", "te\x034,3st"},
	{"te$c[green]4st", "te\x03034st"},
	{"te$c[red,green]9st", "te\x034,039st"},
}

var escapetests = []testcase{
	{"te$c[]st", "te\x03st"},
	{"test$c[]", "test\x03"},
}

var unescapetests = []testcase{
	{"te$st", "test"},
}

func TestEscape(t *testing.T) {
	for _, pair := range tests {
		val := Escape(pair.unescaped)

		if val != pair.escaped {
			t.Error(
				"For", pair.unescaped,
				"expected", pair.escaped,
				"got", val,
			)
		}
	}
	for _, pair := range escapetests {
		val := Escape(pair.unescaped)

		if val != pair.escaped {
			t.Error(
				"For", pair.unescaped,
				"expected", pair.escaped,
				"got", val,
			)
		}
	}
}

func TestUnescape(t *testing.T) {
	for _, pair := range tests {
		val := Unescape(pair.escaped)

		if val != pair.unescaped {
			t.Error(
				"For", pair.escaped,
				"expected", pair.unescaped,
				"got", val,
			)
		}
	}
	for _, pair := range unescapetests {
		val := Unescape(pair.escaped)

		if val != pair.unescaped {
			t.Error(
				"For", pair.escaped,
				"expected", pair.unescaped,
				"got", val,
			)
		}
	}
}
<commit_msg>ircfmt: Add test for the new  raw parsing code<commit_after>package ircfmt

import "testing"

type testcase struct {
	escaped   string
	unescaped string
}

var tests = []testcase{
	{"te$bst", "te\x02st"},
	{"te$c[green]st", "te\x033st"},
	{"te$c[red,green]st", "te\x034,3st"},
	{"te$c[green]4st", "te\x03034st"},
	{"te$c[red,green]9st", "te\x034,039st"},
}

var escapetests = []testcase{
	{"te$c[]st", "te\x03st"},
	{"test$c[]", "test\x03"},
}

var unescapetests = []testcase{
	{"te$st", "test"},
	{"test$c", "test\x03"},
}

func TestEscape(t *testing.T) {
	for _, pair := range tests {
		val := Escape(pair.unescaped)

		if val != pair.escaped {
			t.Error(
				"For", pair.unescaped,
				"expected", pair.escaped,
				"got", val,
			)
		}
	}
	for _, pair := range escapetests {
		val := Escape(pair.unescaped)

		if val != pair.escaped {
			t.Error(
				"For", pair.unescaped,
				"expected", pair.escaped,
				"got", val,
			)
		}
	}
}

func TestUnescape(t *testing.T) {
	for _, pair := range tests {
		val := Unescape(pair.escaped)

		if val != pair.unescaped {
			t.Error(
				"For", pair.escaped,
				"expected", pair.unescaped,
				"got", val,
			)
		}
	}
	for _, pair := range unescapetests {
		val := Unescape(pair.escaped)

		if val != pair.unescaped {
			t.Error(
				"For", pair.escaped,
				"expected", pair.unescaped,
				"got", val,
			)
		}
	}
}
<|endoftext|>


<commit_before>package irmaclient

import (
	"encoding/json"
	"encoding/xml"
	"html"
	"io/ioutil"
	"math/big"
	"time"

	"github.com/go-errors/errors"
	"github.com/mhe/gabi"
	"github.com/privacybydesign/irmago"
	"github.com/privacybydesign/irmago/internal/fs"
)

// This file contains the update mechanism for Client
// as well as updates themselves.

type update struct {
	When    irma.Timestamp
	Number  int
	Success bool
	Error   *string
}

var clientUpdates = []func(client *Client) error{
	// Convert old cardemu.xml Android storage to our own storage format
	func(client *Client) error {
		_, err := client.ParseAndroidStorage()
		return err
	},

	// Adding scheme manager index, signature and public key
	// Check the signatures of all scheme managers, if any is not ok,
	// copy the entire irma_configuration folder from assets
	func(client *Client) error {
		conf := client.Configuration
		if len(conf.DisabledSchemeManagers) > 0 {
			return conf.CopyFromAssets(true)
		}
		for manager := range conf.SchemeManagers {
			valid, err := conf.VerifySignature(manager)
			if err != nil || !valid {
				return conf.CopyFromAssets(true)
			}
		}
		return nil
	},

	// Rename config -> preferences
	func(client *Client) (err error) {
		exists, err := fs.PathExists(client.storage.path("config"))
		if !exists || err != nil {
			return
		}
		oldStruct := &struct {
			SendCrashReports bool
		}{}
		// Load old file, convert to new struct, and save
		err = client.storage.load(oldStruct, "config")
		if err != nil {
			return err
		}
		client.Preferences = Preferences{
			EnableCrashReporting: oldStruct.SendCrashReports,
		}
		return client.storage.StorePreferences(client.Preferences)
	},
}

// update performs any function from clientUpdates that has not
// already been executed in the past, keeping track of previously executed updates
// in the file at updatesFile.
func (client *Client) update() error {
	// Load and parse file containing info about already performed updates
	var err error
	if client.updates, err = client.storage.LoadUpdates(); err != nil {
		return err
	}

	// Perform all new updates
	for i := len(client.updates); i < len(clientUpdates); i++ {
		if clientUpdates[i] == nil {
			continue
		}
		err = clientUpdates[i](client)
		u := update{
			When:    irma.Timestamp(time.Now()),
			Number:  i,
			Success: err == nil,
		}
		if err != nil {
			str := err.Error()
			u.Error = &str
		}
		client.updates = append(client.updates, u)
		if err != nil {
			break
		}
	}

	storeErr := client.storage.StoreUpdates(client.updates)
	if storeErr != nil {
		return storeErr
	}
	return err
}

// ParseAndroidStorage parses an Android cardemu.xml shared preferences file
// from the old Android IRMA app, parsing its credentials into the current instance,
// and saving them to storage.
// CAREFUL: this method overwrites any existing secret keys and attributes on storage.
func (client *Client) ParseAndroidStorage() (present bool, err error) {
	if client.androidStoragePath == "" {
		return false, nil
	}

	cardemuXML := client.androidStoragePath + "/shared_prefs/cardemu.xml"
	present, err = fs.PathExists(cardemuXML)
	if err != nil || !present {
		return
	}
	present = true

	bytes, err := ioutil.ReadFile(cardemuXML)
	if err != nil {
		return
	}
	parsedxml := struct {
		Strings []struct {
			Name    string `xml:"name,attr"`
			Content string `xml:",chardata"`
		} `xml:"string"`
	}{}
	if err = xml.Unmarshal(bytes, &parsedxml); err != nil {
		return
	}

	parsedjson := make(map[string][]*struct {
		Signature    *gabi.CLSignature `json:"signature"`
		Pk           *gabi.PublicKey   `json:"-"`
		Attributes   []*big.Int        `json:"attributes"`
		SharedPoints []*big.Int        `json:"public_sks"`
	})
	client.keyshareServers = make(map[irma.SchemeManagerIdentifier]*keyshareServer)
	for _, xmltag := range parsedxml.Strings {
		if xmltag.Name == "credentials" {
			jsontag := html.UnescapeString(xmltag.Content)
			if err = json.Unmarshal([]byte(jsontag), &parsedjson); err != nil {
				return
			}
		}
		if xmltag.Name == "keyshare" {
			jsontag := html.UnescapeString(xmltag.Content)
			if err = json.Unmarshal([]byte(jsontag), &client.keyshareServers); err != nil {
				return
			}
		}
		if xmltag.Name == "KeyshareKeypairs" {
			jsontag := html.UnescapeString(xmltag.Content)
			keys := make([]*paillierPrivateKey, 0, 3)
			if err = json.Unmarshal([]byte(jsontag), &keys); err != nil {
				return
			}
			client.paillierKeyCache = keys[0]
		}
	}

	for _, list := range parsedjson {
		client.secretkey = &secretKey{Key: list[0].Attributes[0]}
		for _, oldcred := range list {
			gabicred := &gabi.Credential{
				Attributes: oldcred.Attributes,
				Signature:  oldcred.Signature,
			}
			if oldcred.SharedPoints != nil && len(oldcred.SharedPoints) > 0 {
				gabicred.Signature.KeyshareP = oldcred.SharedPoints[0]
			}
			var cred *credential
			if cred, err = newCredential(gabicred, client.Configuration); err != nil {
				return
			}
			if cred.CredentialType() == nil {
				err = errors.New("cannot add unknown credential type")
				return
			}

			if err = client.addCredential(cred, false); err != nil {
				return
			}
		}
	}

	if len(client.credentials) > 0 {
		if err = client.storage.StoreAttributes(client.attributes); err != nil {
			return
		}
		if err = client.storage.StoreSecretKey(client.secretkey); err != nil {
			return
		}
	}

	if len(client.keyshareServers) > 0 {
		if err = client.storage.StoreKeyshareServers(client.keyshareServers); err != nil {
			return
		}
	}
	client.UnenrolledSchemeManagers = client.unenrolledSchemeManagers()

	if err = client.storage.StorePaillierKeys(client.paillierKeyCache); err != nil {
		return
	}
	if client.paillierKeyCache == nil {
		client.paillierKey(false) // trigger calculating a new one
	}
	return
}
<commit_msg>Add update function to copy irma_configuration out of assets<commit_after>package irmaclient

import (
	"encoding/json"
	"encoding/xml"
	"html"
	"io/ioutil"
	"math/big"
	"time"

	"github.com/go-errors/errors"
	"github.com/mhe/gabi"
	"github.com/privacybydesign/irmago"
	"github.com/privacybydesign/irmago/internal/fs"
)

// This file contains the update mechanism for Client
// as well as updates themselves.

type update struct {
	When    irma.Timestamp
	Number  int
	Success bool
	Error   *string
}

var clientUpdates = []func(client *Client) error{
	// Convert old cardemu.xml Android storage to our own storage format
	func(client *Client) error {
		_, err := client.ParseAndroidStorage()
		return err
	},

	// Adding scheme manager index, signature and public key
	// Check the signatures of all scheme managers, if any is not ok,
	// copy the entire irma_configuration folder from assets
	func(client *Client) error {
		conf := client.Configuration
		if len(conf.DisabledSchemeManagers) > 0 {
			return conf.CopyFromAssets(true)
		}
		for manager := range conf.SchemeManagers {
			valid, err := conf.VerifySignature(manager)
			if err != nil || !valid {
				return conf.CopyFromAssets(true)
			}
		}
		return nil
	},

	// Rename config -> preferences
	func(client *Client) (err error) {
		exists, err := fs.PathExists(client.storage.path("config"))
		if !exists || err != nil {
			return
		}
		oldStruct := &struct {
			SendCrashReports bool
		}{}
		// Load old file, convert to new struct, and save
		err = client.storage.load(oldStruct, "config")
		if err != nil {
			return err
		}
		client.Preferences = Preferences{
			EnableCrashReporting: oldStruct.SendCrashReports,
		}
		return client.storage.StorePreferences(client.Preferences)
	},

	// Copy new irma_configuration out of assets
	func(client *Client) (err error) {
		return client.Configuration.CopyFromAssets(true)
	},
}

// update performs any function from clientUpdates that has not
// already been executed in the past, keeping track of previously executed updates
// in the file at updatesFile.
func (client *Client) update() error {
	// Load and parse file containing info about already performed updates
	var err error
	if client.updates, err = client.storage.LoadUpdates(); err != nil {
		return err
	}

	// Perform all new updates
	for i := len(client.updates); i < len(clientUpdates); i++ {
		if clientUpdates[i] == nil {
			continue
		}
		err = clientUpdates[i](client)
		u := update{
			When:    irma.Timestamp(time.Now()),
			Number:  i,
			Success: err == nil,
		}
		if err != nil {
			str := err.Error()
			u.Error = &str
		}
		client.updates = append(client.updates, u)
		if err != nil {
			break
		}
	}

	storeErr := client.storage.StoreUpdates(client.updates)
	if storeErr != nil {
		return storeErr
	}
	return err
}

// ParseAndroidStorage parses an Android cardemu.xml shared preferences file
// from the old Android IRMA app, parsing its credentials into the current instance,
// and saving them to storage.
// CAREFUL: this method overwrites any existing secret keys and attributes on storage.
func (client *Client) ParseAndroidStorage() (present bool, err error) {
	if client.androidStoragePath == "" {
		return false, nil
	}

	cardemuXML := client.androidStoragePath + "/shared_prefs/cardemu.xml"
	present, err = fs.PathExists(cardemuXML)
	if err != nil || !present {
		return
	}
	present = true

	bytes, err := ioutil.ReadFile(cardemuXML)
	if err != nil {
		return
	}
	parsedxml := struct {
		Strings []struct {
			Name    string `xml:"name,attr"`
			Content string `xml:",chardata"`
		} `xml:"string"`
	}{}
	if err = xml.Unmarshal(bytes, &parsedxml); err != nil {
		return
	}

	parsedjson := make(map[string][]*struct {
		Signature    *gabi.CLSignature `json:"signature"`
		Pk           *gabi.PublicKey   `json:"-"`
		Attributes   []*big.Int        `json:"attributes"`
		SharedPoints []*big.Int        `json:"public_sks"`
	})
	client.keyshareServers = make(map[irma.SchemeManagerIdentifier]*keyshareServer)
	for _, xmltag := range parsedxml.Strings {
		if xmltag.Name == "credentials" {
			jsontag := html.UnescapeString(xmltag.Content)
			if err = json.Unmarshal([]byte(jsontag), &parsedjson); err != nil {
				return
			}
		}
		if xmltag.Name == "keyshare" {
			jsontag := html.UnescapeString(xmltag.Content)
			if err = json.Unmarshal([]byte(jsontag), &client.keyshareServers); err != nil {
				return
			}
		}
		if xmltag.Name == "KeyshareKeypairs" {
			jsontag := html.UnescapeString(xmltag.Content)
			keys := make([]*paillierPrivateKey, 0, 3)
			if err = json.Unmarshal([]byte(jsontag), &keys); err != nil {
				return
			}
			client.paillierKeyCache = keys[0]
		}
	}

	for _, list := range parsedjson {
		client.secretkey = &secretKey{Key: list[0].Attributes[0]}
		for _, oldcred := range list {
			gabicred := &gabi.Credential{
				Attributes: oldcred.Attributes,
				Signature:  oldcred.Signature,
			}
			if oldcred.SharedPoints != nil && len(oldcred.SharedPoints) > 0 {
				gabicred.Signature.KeyshareP = oldcred.SharedPoints[0]
			}
			var cred *credential
			if cred, err = newCredential(gabicred, client.Configuration); err != nil {
				return
			}
			if cred.CredentialType() == nil {
				err = errors.New("cannot add unknown credential type")
				return
			}

			if err = client.addCredential(cred, false); err != nil {
				return
			}
		}
	}

	if len(client.credentials) > 0 {
		if err = client.storage.StoreAttributes(client.attributes); err != nil {
			return
		}
		if err = client.storage.StoreSecretKey(client.secretkey); err != nil {
			return
		}
	}

	if len(client.keyshareServers) > 0 {
		if err = client.storage.StoreKeyshareServers(client.keyshareServers); err != nil {
			return
		}
	}
	client.UnenrolledSchemeManagers = client.unenrolledSchemeManagers()

	if err = client.storage.StorePaillierKeys(client.paillierKeyCache); err != nil {
		return
	}
	if client.paillierKeyCache == nil {
		client.paillierKey(false) // trigger calculating a new one
	}
	return
}
<|endoftext|>


<commit_before>// Copyright 2016 Google Inc. All rights reserved.
//
// 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 query

import (
	"log"
	"reflect"
	"regexp/syntax"
	"strings"
	"testing"
)

var opnames = map[syntax.Op]string{
	syntax.OpNoMatch:        "OpNoMatch",
	syntax.OpEmptyMatch:     "OpEmptyMatch",
	syntax.OpLiteral:        "OpLiteral",
	syntax.OpCharClass:      "OpCharClass",
	syntax.OpAnyCharNotNL:   "OpAnyCharNotNL",
	syntax.OpAnyChar:        "OpAnyChar",
	syntax.OpBeginLine:      "OpBeginLine",
	syntax.OpEndLine:        "OpEndLine",
	syntax.OpBeginText:      "OpBeginText",
	syntax.OpEndText:        "OpEndText",
	syntax.OpWordBoundary:   "OpWordBoundary",
	syntax.OpNoWordBoundary: "OpNoWordBoundary",
	syntax.OpCapture:        "OpCapture",
	syntax.OpStar:           "OpStar",
	syntax.OpPlus:           "OpPlus",
	syntax.OpQuest:          "OpQuest",
	syntax.OpRepeat:         "OpRepeat",
	syntax.OpConcat:         "OpConcat",
	syntax.OpAlternate:      "OpAlternate",
}

func printRegexp(r *syntax.Regexp, lvl int) {
	log.Printf("%s%s ch: %d", strings.Repeat(" ", lvl), opnames[r.Op], len(r.Sub))
	for _, s := range r.Sub {
		printRegexp(s, lvl+1)
	}
}

func TestRegexpParse(t *testing.T) {
	type testcase struct {
		in   string
		want Q
	}

	cases := []testcase{
		{"(foo|)bar", &Substring{Pattern: "bar"}},
		{"(foo|)", &Const{true}},
		{"(foo|bar)baz.*bla", &And{[]Q{
			&Or{[]Q{
				&Substring{Pattern: "foo"},
				&Substring{Pattern: "bar"},
			}},
			&Substring{Pattern: "baz"},
			&Substring{Pattern: "bla"},
		}}},
		{"^[a-z](People)+barrabas$",
			&And{[]Q{
				&Substring{Pattern: "People"},
				&Substring{Pattern: "barrabas"},
			}}},
	}

	for _, c := range cases {
		r, err := syntax.Parse(c.in, syntax.Perl)
		if err != nil {
			t.Errorf("Parse(%q): %v", c.in, err)
			continue
		}

		got := RegexpToQuery(r, 3)
		if !reflect.DeepEqual(c.want, got) {
			t.Errorf("regexpToQuery(%q): got %v, want %v", c.in, got, c.want)
		}
	}
}

func TestLowerRegexp(t *testing.T) {
	in := "[a-zA-Z]fooBAR"
	re := mustParseRE(in)
	in = re.String()
	got := LowerRegexp(re)
	want := "[a-za-z]foobar"
	if got.String() != want {
		t.Errorf("got %s, want %s", got, want)
	}

	if re.String() != in {
		t.Errorf("got mutated original %s want %s", re.String(), in)
	}
}
<commit_msg>query: Always use printRegexp in tests<commit_after>// Copyright 2016 Google Inc. All rights reserved.
//
// 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 query

import (
	"reflect"
	"regexp/syntax"
	"strings"
	"testing"
)

var opnames = map[syntax.Op]string{
	syntax.OpNoMatch:        "OpNoMatch",
	syntax.OpEmptyMatch:     "OpEmptyMatch",
	syntax.OpLiteral:        "OpLiteral",
	syntax.OpCharClass:      "OpCharClass",
	syntax.OpAnyCharNotNL:   "OpAnyCharNotNL",
	syntax.OpAnyChar:        "OpAnyChar",
	syntax.OpBeginLine:      "OpBeginLine",
	syntax.OpEndLine:        "OpEndLine",
	syntax.OpBeginText:      "OpBeginText",
	syntax.OpEndText:        "OpEndText",
	syntax.OpWordBoundary:   "OpWordBoundary",
	syntax.OpNoWordBoundary: "OpNoWordBoundary",
	syntax.OpCapture:        "OpCapture",
	syntax.OpStar:           "OpStar",
	syntax.OpPlus:           "OpPlus",
	syntax.OpQuest:          "OpQuest",
	syntax.OpRepeat:         "OpRepeat",
	syntax.OpConcat:         "OpConcat",
	syntax.OpAlternate:      "OpAlternate",
}

func printRegexp(t *testing.T, r *syntax.Regexp, lvl int) {
	t.Logf("%s%s ch: %d", strings.Repeat(" ", lvl), opnames[r.Op], len(r.Sub))
	for _, s := range r.Sub {
		printRegexp(t, s, lvl+1)
	}
}

func TestRegexpParse(t *testing.T) {
	type testcase struct {
		in   string
		want Q
	}

	cases := []testcase{
		{"(foo|)bar", &Substring{Pattern: "bar"}},
		{"(foo|)", &Const{true}},
		{"(foo|bar)baz.*bla", &And{[]Q{
			&Or{[]Q{
				&Substring{Pattern: "foo"},
				&Substring{Pattern: "bar"},
			}},
			&Substring{Pattern: "baz"},
			&Substring{Pattern: "bla"},
		}}},
		{"^[a-z](People)+barrabas$",
			&And{[]Q{
				&Substring{Pattern: "People"},
				&Substring{Pattern: "barrabas"},
			}}},
	}

	for _, c := range cases {
		r, err := syntax.Parse(c.in, syntax.Perl)
		if err != nil {
			t.Errorf("Parse(%q): %v", c.in, err)
			continue
		}

		got := RegexpToQuery(r, 3)
		if !reflect.DeepEqual(c.want, got) {
			printRegexp(t, r, 0)
			t.Errorf("regexpToQuery(%q): got %v, want %v", c.in, got, c.want)
		}
	}
}

func TestLowerRegexp(t *testing.T) {
	in := "[a-zA-Z]fooBAR"
	re := mustParseRE(in)
	in = re.String()
	got := LowerRegexp(re)
	want := "[a-za-z]foobar"
	if got.String() != want {
		printRegexp(t, re, 0)
		printRegexp(t, got, 0)
		t.Errorf("got %s, want %s", got, want)
	}

	if re.String() != in {
		t.Errorf("got mutated original %s want %s", re.String(), in)
	}
}
<|endoftext|>


<commit_before>package lib

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"testing"
)

type Mux map[string]func(wr http.ResponseWriter, r *http.Request)

func (m Mux) ToHandler() (h http.Handler) {

	serveMux := http.NewServeMux()
	for p, f := range m {
		serveMux.HandleFunc(p, f)
	}
	return serveMux
}

type MuxHandlers struct {
	auth    Mux
	brain   Mux
	billing Mux
	spp     Mux
	api     Mux
}

func (mh MuxHandlers) ToHandlers() (h Handlers) {
	h.auth = mh.auth.ToHandler()
	if mh.auth == nil {
		h.auth = nil
	}
	h.brain = mh.brain.ToHandler()
	h.billing = mh.billing.ToHandler()
	h.spp = mh.spp.ToHandler()
	h.api = mh.api.ToHandler()
	return
}

type Handlers struct {
	auth    http.Handler
	brain   http.Handler
	billing http.Handler
	spp     http.Handler
	api     http.Handler
}

func (h *Handlers) Fill(t *testing.T) {
	if h.brain == nil {
		h.brain = mkNilHandler(t)
	}
	if h.billing == nil {
		h.billing = mkNilHandler(t)
	}
	if h.spp == nil {
		h.spp = mkNilHandler(t)
	}
	if h.api == nil {
		h.api = mkNilHandler(t)
	}
}

type Servers struct {
	auth    *httptest.Server
	brain   *httptest.Server
	billing *httptest.Server
	spp     *httptest.Server
	api     *httptest.Server
}

func (s *Servers) Close() {
	if s.auth != nil {
		s.auth.Close()
	}
	if s.brain != nil {
		s.brain.Close()
	}
	if s.billing != nil {
		s.billing.Close()
	}
	if s.spp != nil {
		s.spp.Close()
	}
	if s.api != nil {
		s.api.Close()
	}
	s.auth = nil
	s.brain = nil
	s.billing = nil
	s.spp = nil
	s.api = nil
}

// URLs creates an EndpointURLs
func (s Servers) URLs() (urls EndpointURLs) {
	urls.API = s.api.URL
	urls.Auth = s.auth.URL
	urls.Billing = s.billing.URL
	urls.Brain = s.brain.URL
	urls.SPP = s.spp.URL
	return
}

// Client makes a bytemarkClient for these Servers
func (s Servers) Client() (c Client, err error) {
	urls := s.URLs()
	c, err = NewWithURLs(urls)
	if err != nil {
		return nil, err
	}
	c.AllowInsecureRequests()
	return
}

// mkTestClientAndServers constructs httptest Servers for a pretend auth and API endpoint, then constructs a Client that uses those servers.
// Used to test that the right URLs are being requested and such.
func mkTestClientAndServers(t *testing.T, h Handlers) (c Client, s Servers, err error) {
	h.Fill(t)

	if h.auth != nil {
		s.auth = mkTestServer(h.auth)
	} else {
		s.auth = mkTestAuthServer()
	}
	s.brain = mkTestServer(h.brain)
	s.billing = mkTestServer(h.billing)
	s.api = mkTestServer(h.api)
	s.spp = mkTestServer(h.spp)

	c, err = s.Client()
	return
}

// mkTestServer creates an httptest.Server for the given http.Handler. It's basically an alias for httptest.NewServer. Why did I write it?
func mkTestServer(handler http.Handler) *httptest.Server {
	return httptest.NewServer(handler)
}

func mkNilHandler(t *testing.T) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		t.Fatalf("Unexpected request to a nil server\r\n%s %s", r.Method, r.URL.String())
	})
}

func mkTestAuthServer() *httptest.Server {
	return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprintln(w,
			`{
    "token": "working-auth-token",
    "username": "account",
    "factors": []
}`)
	}))

}

// JSON unmarshals the contents of r.Body into obj.
func unmarshalRequestObject(r *http.Request, obj interface{}) (err error) {
	body, err := ioutil.ReadAll(r.Body)
	if err != nil {
		return
	}
	err = json.Unmarshal(body, obj)
	return
}
<commit_msg>Remove unmarshalRequestObject<commit_after>package lib

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"testing"
)

type Mux map[string]func(wr http.ResponseWriter, r *http.Request)

func (m Mux) ToHandler() (h http.Handler) {

	serveMux := http.NewServeMux()
	for p, f := range m {
		serveMux.HandleFunc(p, f)
	}
	return serveMux
}

type MuxHandlers struct {
	auth    Mux
	brain   Mux
	billing Mux
	spp     Mux
	api     Mux
}

func (mh MuxHandlers) ToHandlers() (h Handlers) {
	h.auth = mh.auth.ToHandler()
	if mh.auth == nil {
		h.auth = nil
	}
	h.brain = mh.brain.ToHandler()
	h.billing = mh.billing.ToHandler()
	h.spp = mh.spp.ToHandler()
	h.api = mh.api.ToHandler()
	return
}

type Handlers struct {
	auth    http.Handler
	brain   http.Handler
	billing http.Handler
	spp     http.Handler
	api     http.Handler
}

func (h *Handlers) Fill(t *testing.T) {
	if h.brain == nil {
		h.brain = mkNilHandler(t)
	}
	if h.billing == nil {
		h.billing = mkNilHandler(t)
	}
	if h.spp == nil {
		h.spp = mkNilHandler(t)
	}
	if h.api == nil {
		h.api = mkNilHandler(t)
	}
}

type Servers struct {
	auth    *httptest.Server
	brain   *httptest.Server
	billing *httptest.Server
	spp     *httptest.Server
	api     *httptest.Server
}

func (s *Servers) Close() {
	if s.auth != nil {
		s.auth.Close()
	}
	if s.brain != nil {
		s.brain.Close()
	}
	if s.billing != nil {
		s.billing.Close()
	}
	if s.spp != nil {
		s.spp.Close()
	}
	if s.api != nil {
		s.api.Close()
	}
	s.auth = nil
	s.brain = nil
	s.billing = nil
	s.spp = nil
	s.api = nil
}

// URLs creates an EndpointURLs
func (s Servers) URLs() (urls EndpointURLs) {
	urls.API = s.api.URL
	urls.Auth = s.auth.URL
	urls.Billing = s.billing.URL
	urls.Brain = s.brain.URL
	urls.SPP = s.spp.URL
	return
}

// Client makes a bytemarkClient for these Servers
func (s Servers) Client() (c Client, err error) {
	urls := s.URLs()
	c, err = NewWithURLs(urls)
	if err != nil {
		return nil, err
	}
	c.AllowInsecureRequests()
	return
}

// mkTestClientAndServers constructs httptest Servers for a pretend auth and API endpoint, then constructs a Client that uses those servers.
// Used to test that the right URLs are being requested and such.
func mkTestClientAndServers(t *testing.T, h Handlers) (c Client, s Servers, err error) {
	h.Fill(t)

	if h.auth != nil {
		s.auth = mkTestServer(h.auth)
	} else {
		s.auth = mkTestAuthServer()
	}
	s.brain = mkTestServer(h.brain)
	s.billing = mkTestServer(h.billing)
	s.api = mkTestServer(h.api)
	s.spp = mkTestServer(h.spp)

	c, err = s.Client()
	return
}

// mkTestServer creates an httptest.Server for the given http.Handler. It's basically an alias for httptest.NewServer. Why did I write it?
func mkTestServer(handler http.Handler) *httptest.Server {
	return httptest.NewServer(handler)
}

func mkNilHandler(t *testing.T) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		t.Fatalf("Unexpected request to a nil server\r\n%s %s", r.Method, r.URL.String())
	})
}

func mkTestAuthServer() *httptest.Server {
	return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprintln(w,
			`{
    "token": "working-auth-token",
    "username": "account",
    "factors": []
}`)
	}))

}
<|endoftext|>