<commit_before>package display

import (
	"bytes"
	"fmt"
	"path"
	"regexp"
	"strconv"
	"strings"
	"unicode/utf8"

	luar "layeh.com/gopher-luar"

	runewidth "github.com/mattn/go-runewidth"
	lua "github.com/yuin/gopher-lua"
	"github.com/zyedidia/micro/internal/buffer"
	"github.com/zyedidia/micro/internal/config"
	ulua "github.com/zyedidia/micro/internal/lua"
	"github.com/zyedidia/micro/internal/screen"
	"github.com/zyedidia/micro/internal/util"
)

// StatusLine represents the information line at the bottom
// of each window
// It gives information such as filename, whether the file has been
// modified, filetype, cursor location
type StatusLine struct {
	Info map[string]func(*buffer.Buffer) string

	win *BufWindow
}

var statusInfo = map[string]func(*buffer.Buffer) string{
	"filename": func(b *buffer.Buffer) string {
		if b.Settings["basename"].(bool) {
			return path.Base(b.GetName())
		}
		return b.GetName()
	},
	"line": func(b *buffer.Buffer) string {
		return strconv.Itoa(b.GetActiveCursor().Y + 1)
	},
	"col": func(b *buffer.Buffer) string {
		return strconv.Itoa(b.GetActiveCursor().X + 1)
	},
	"modified": func(b *buffer.Buffer) string {
		if b.Modified() {
			return "+ "
		}
		return ""
	},
}

func SetStatusInfoFnLua(s string, fn string) {
	luaFn := strings.Split(fn, ".")
	if len(luaFn) <= 1 {
		return
	}
	plName, plFn := luaFn[0], luaFn[1]
	pl := config.FindPlugin(plName)
	if pl == nil {
		return
	}
	statusInfo[s] = func(b *buffer.Buffer) string {
		if pl == nil || !pl.IsEnabled() {
			return ""
		}
		val, err := pl.Call(plFn, luar.New(ulua.L, b))
		if err == nil {
			if v, ok := val.(lua.LString); !ok {
				screen.TermMessage(plFn, "should return a string")
				return ""
			} else {
				return string(v)
			}
		}
		return ""
	}
}

// NewStatusLine returns a statusline bound to a window
func NewStatusLine(win *BufWindow) *StatusLine {
	s := new(StatusLine)
	s.win = win
	return s
}

// FindOpt finds a given option in the current buffer's settings
func (s *StatusLine) FindOpt(opt string) interface{} {
	if val, ok := s.win.Buf.Settings[opt]; ok {
		return val
	}
	return "null"
}

var formatParser = regexp.MustCompile(`\$\(.+?\)`)

// Display draws the statusline to the screen
func (s *StatusLine) Display() {
	// We'll draw the line at the lowest line in the window
	y := s.win.Height + s.win.Y - 1

	b := s.win.Buf
	// autocomplete suggestions (for the buffer, not for the infowindow)
	if b.HasSuggestions && len(b.Suggestions) > 1 {
		statusLineStyle := config.DefStyle.Reverse(true)
		if style, ok := config.Colorscheme["statusline"]; ok {
			statusLineStyle = style
		}
		keymenuOffset := 0
		if config.GetGlobalOption("keymenu").(bool) {
			keymenuOffset = len(keydisplay)
		}
		x := 0
		for j, sug := range b.Suggestions {
			style := statusLineStyle
			if b.CurSuggestion == j {
				style = style.Reverse(true)
			}
			for _, r := range sug {
				screen.Screen.SetContent(x, y-keymenuOffset, r, nil, style)
				x++
				if x >= s.win.Width {
					return
				}
			}
			screen.Screen.SetContent(x, y-keymenuOffset, ' ', nil, statusLineStyle)
			x++
			if x >= s.win.Width {
				return
			}
		}

		for x < s.win.Width {
			screen.Screen.SetContent(x, y-keymenuOffset, ' ', nil, statusLineStyle)
			x++
		}
		return
	}

	formatter := func(match []byte) []byte {
		name := match[2 : len(match)-1]
		if bytes.HasPrefix(name, []byte("opt")) {
			option := name[4:]
			return []byte(fmt.Sprint(s.FindOpt(string(option))))
		} else if bytes.HasPrefix(name, []byte("bind")) {
			binding := string(name[5:])
			for k, v := range config.Bindings {
				if v == binding {
					return []byte(k)
				}
			}
			return []byte("null")
		} else {
			if fn, ok := statusInfo[string(name)]; ok {
				return []byte(fn(s.win.Buf))
			}
			return []byte{}
		}
	}

	leftText := []byte(s.win.Buf.Settings["statusformatl"].(string))
	leftText = formatParser.ReplaceAllFunc(leftText, formatter)
	rightText := []byte(s.win.Buf.Settings["statusformatr"].(string))
	rightText = formatParser.ReplaceAllFunc(rightText, formatter)

	statusLineStyle := config.DefStyle.Reverse(true)
	if style, ok := config.Colorscheme["statusline"]; ok {
		statusLineStyle = style
	}

	leftLen := util.StringWidth(leftText, utf8.RuneCount(leftText), 1)
	rightLen := util.StringWidth(rightText, utf8.RuneCount(rightText), 1)

	winX := s.win.X
	for x := 0; x < s.win.Width; x++ {
		if x < leftLen {
			r, size := utf8.DecodeRune(leftText)
			leftText = leftText[size:]
			rw := runewidth.RuneWidth(r)
			for j := 0; j < rw; j++ {
				c := r
				if j > 0 {
					c = ' '
					x++
				}
				screen.Screen.SetContent(winX+x, y, c, nil, statusLineStyle)
			}
		} else if x >= s.win.Width-rightLen && x < rightLen+s.win.Width-rightLen {
			r, size := utf8.DecodeRune(rightText)
			rightText = rightText[size:]
			rw := runewidth.RuneWidth(r)
			for j := 0; j < rw; j++ {
				c := r
				if j > 0 {
					c = ' '
					x++
				}
				screen.Screen.SetContent(winX+x, y, c, nil, statusLineStyle)
			}
		} else {
			screen.Screen.SetContent(winX+x, y, ' ', nil, statusLineStyle)
		}
	}
}
<commit_msg>Improve lua interface for statusline<commit_after>package display

import (
	"bytes"
	"fmt"
	"path"
	"regexp"
	"strconv"
	"strings"
	"unicode/utf8"

	luar "layeh.com/gopher-luar"

	runewidth "github.com/mattn/go-runewidth"
	lua "github.com/yuin/gopher-lua"
	"github.com/zyedidia/micro/internal/buffer"
	"github.com/zyedidia/micro/internal/config"
	ulua "github.com/zyedidia/micro/internal/lua"
	"github.com/zyedidia/micro/internal/screen"
	"github.com/zyedidia/micro/internal/util"
)

// StatusLine represents the information line at the bottom
// of each window
// It gives information such as filename, whether the file has been
// modified, filetype, cursor location
type StatusLine struct {
	Info map[string]func(*buffer.Buffer) string

	win *BufWindow
}

var statusInfo = map[string]func(*buffer.Buffer) string{
	"filename": func(b *buffer.Buffer) string {
		if b.Settings["basename"].(bool) {
			return path.Base(b.GetName())
		}
		return b.GetName()
	},
	"line": func(b *buffer.Buffer) string {
		return strconv.Itoa(b.GetActiveCursor().Y + 1)
	},
	"col": func(b *buffer.Buffer) string {
		return strconv.Itoa(b.GetActiveCursor().X + 1)
	},
	"modified": func(b *buffer.Buffer) string {
		if b.Modified() {
			return "+ "
		}
		return ""
	},
}

func SetStatusInfoFnLua(fn string) {
	luaFn := strings.Split(fn, ".")
	if len(luaFn) <= 1 {
		return
	}
	plName, plFn := luaFn[0], luaFn[1]
	pl := config.FindPlugin(plName)
	if pl == nil {
		return
	}
	statusInfo[fn] = func(b *buffer.Buffer) string {
		if pl == nil || !pl.IsEnabled() {
			return ""
		}
		val, err := pl.Call(plFn, luar.New(ulua.L, b))
		if err == nil {
			if v, ok := val.(lua.LString); !ok {
				screen.TermMessage(plFn, "should return a string")
				return ""
			} else {
				return string(v)
			}
		}
		return ""
	}
}

// NewStatusLine returns a statusline bound to a window
func NewStatusLine(win *BufWindow) *StatusLine {
	s := new(StatusLine)
	s.win = win
	return s
}

// FindOpt finds a given option in the current buffer's settings
func (s *StatusLine) FindOpt(opt string) interface{} {
	if val, ok := s.win.Buf.Settings[opt]; ok {
		return val
	}
	return "null"
}

var formatParser = regexp.MustCompile(`\$\(.+?\)`)

// Display draws the statusline to the screen
func (s *StatusLine) Display() {
	// We'll draw the line at the lowest line in the window
	y := s.win.Height + s.win.Y - 1

	b := s.win.Buf
	// autocomplete suggestions (for the buffer, not for the infowindow)
	if b.HasSuggestions && len(b.Suggestions) > 1 {
		statusLineStyle := config.DefStyle.Reverse(true)
		if style, ok := config.Colorscheme["statusline"]; ok {
			statusLineStyle = style
		}
		keymenuOffset := 0
		if config.GetGlobalOption("keymenu").(bool) {
			keymenuOffset = len(keydisplay)
		}
		x := 0
		for j, sug := range b.Suggestions {
			style := statusLineStyle
			if b.CurSuggestion == j {
				style = style.Reverse(true)
			}
			for _, r := range sug {
				screen.Screen.SetContent(x, y-keymenuOffset, r, nil, style)
				x++
				if x >= s.win.Width {
					return
				}
			}
			screen.Screen.SetContent(x, y-keymenuOffset, ' ', nil, statusLineStyle)
			x++
			if x >= s.win.Width {
				return
			}
		}

		for x < s.win.Width {
			screen.Screen.SetContent(x, y-keymenuOffset, ' ', nil, statusLineStyle)
			x++
		}
		return
	}

	formatter := func(match []byte) []byte {
		name := match[2 : len(match)-1]
		if bytes.HasPrefix(name, []byte("opt")) {
			option := name[4:]
			return []byte(fmt.Sprint(s.FindOpt(string(option))))
		} else if bytes.HasPrefix(name, []byte("bind")) {
			binding := string(name[5:])
			for k, v := range config.Bindings {
				if v == binding {
					return []byte(k)
				}
			}
			return []byte("null")
		} else {
			if fn, ok := statusInfo[string(name)]; ok {
				return []byte(fn(s.win.Buf))
			}
			return []byte{}
		}
	}

	leftText := []byte(s.win.Buf.Settings["statusformatl"].(string))
	leftText = formatParser.ReplaceAllFunc(leftText, formatter)
	rightText := []byte(s.win.Buf.Settings["statusformatr"].(string))
	rightText = formatParser.ReplaceAllFunc(rightText, formatter)

	statusLineStyle := config.DefStyle.Reverse(true)
	if style, ok := config.Colorscheme["statusline"]; ok {
		statusLineStyle = style
	}

	leftLen := util.StringWidth(leftText, utf8.RuneCount(leftText), 1)
	rightLen := util.StringWidth(rightText, utf8.RuneCount(rightText), 1)

	winX := s.win.X
	for x := 0; x < s.win.Width; x++ {
		if x < leftLen {
			r, size := utf8.DecodeRune(leftText)
			leftText = leftText[size:]
			rw := runewidth.RuneWidth(r)
			for j := 0; j < rw; j++ {
				c := r
				if j > 0 {
					c = ' '
					x++
				}
				screen.Screen.SetContent(winX+x, y, c, nil, statusLineStyle)
			}
		} else if x >= s.win.Width-rightLen && x < rightLen+s.win.Width-rightLen {
			r, size := utf8.DecodeRune(rightText)
			rightText = rightText[size:]
			rw := runewidth.RuneWidth(r)
			for j := 0; j < rw; j++ {
				c := r
				if j > 0 {
					c = ' '
					x++
				}
				screen.Screen.SetContent(winX+x, y, c, nil, statusLineStyle)
			}
		} else {
			screen.Screen.SetContent(winX+x, y, ' ', nil, statusLineStyle)
		}
	}
}
<|endoftext|>


<commit_before>package gormigrate

import (
	"errors"
	"fmt"

	"github.com/jinzhu/gorm"
)

const (
	initSchemaMigrationID = "SCHEMA_INIT"
)

// MigrateFunc is the func signature for migrating.
type MigrateFunc func(*gorm.DB) error

// RollbackFunc is the func signature for rollbacking.
type RollbackFunc func(*gorm.DB) error

// InitSchemaFunc is the func signature for initializing the schema.
type InitSchemaFunc func(*gorm.DB) error

// Options define options for all migrations.
type Options struct {
	// TableName is the migration table.
	TableName string
	// IDColumnName is the name of column where the migration id will be stored.
	IDColumnName string
	// IDColumnSize is the length of the migration id column
	IDColumnSize int
	// UseTransaction makes Gormigrate execute migrations inside a single transaction.
	// Keep in mind that not all databases support DDL commands inside transactions.
	UseTransaction bool
}

// Migration represents a database migration (a modification to be made on the database).
type Migration struct {
	// ID is the migration identifier. Usually a timestamp like "201601021504".
	ID string
	// Migrate is a function that will br executed while running this migration.
	Migrate MigrateFunc
	// Rollback will be executed on rollback. Can be nil.
	Rollback RollbackFunc
}

// Gormigrate represents a collection of all migrations of a database schema.
type Gormigrate struct {
	db         *gorm.DB
	tx         *gorm.DB
	options    *Options
	migrations []*Migration
	initSchema InitSchemaFunc
}

// ReservedIDError is returned when a migration is using a reserved ID
type ReservedIDError struct {
	ID string
}

func (e *ReservedIDError) Error() string {
	return fmt.Sprintf(`gormigrate: Reserved migration ID: "%s"`, e.ID)
}

// DuplicatedIDError is returned when more than one migration have the same ID
type DuplicatedIDError struct {
	ID string
}

func (e *DuplicatedIDError) Error() string {
	return fmt.Sprintf(`gormigrate: Duplicated migration ID: "%s"`, e.ID)
}

var (
	// DefaultOptions can be used if you don't want to think about options.
	DefaultOptions = &Options{
		TableName:      "migrations",
		IDColumnName:   "id",
		IDColumnSize:   255,
		UseTransaction: false,
	}

	// ErrRollbackImpossible is returned when trying to rollback a migration
	// that has no rollback function.
	ErrRollbackImpossible = errors.New("gormigrate: It's impossible to rollback this migration")

	// ErrNoMigrationDefined is returned when no migration is defined.
	ErrNoMigrationDefined = errors.New("gormigrate: No migration defined")

	// ErrMissingID is returned when the ID od migration is equal to ""
	ErrMissingID = errors.New("gormigrate: Missing ID in migration")

	// ErrNoRunMigration is returned when any run migration was found while
	// running RollbackLast
	ErrNoRunMigration = errors.New("gormigrate: Could not find last run migration")

	// ErrMigrationIDDoesNotExist is returned when migrating or rolling back to a migration ID that
	// does not exist in the list of migrations
	ErrMigrationIDDoesNotExist = errors.New("gormigrate: Tried to migrate to an ID that doesn't exist")
)

// New returns a new Gormigrate.
func New(db *gorm.DB, options *Options, migrations []*Migration) *Gormigrate {
	if options.TableName == "" {
		options.TableName = DefaultOptions.TableName
	}
	if options.IDColumnName == "" {
		options.IDColumnName = DefaultOptions.IDColumnName
	}
	if options.IDColumnSize == 0 {
		options.IDColumnSize = DefaultOptions.IDColumnSize
	}
	return &Gormigrate{
		db:         db,
		options:    options,
		migrations: migrations,
	}
}

// InitSchema sets a function that is run if no migration is found.
// The idea is preventing to run all migrations when a new clean database
// is being migrating. In this function you should create all tables and
// foreign key necessary to your application.
func (g *Gormigrate) InitSchema(initSchema InitSchemaFunc) {
	g.initSchema = initSchema
}

// Migrate executes all migrations that did not run yet.
func (g *Gormigrate) Migrate() error {
	if !g.hasMigrations() {
		return ErrNoMigrationDefined
	}
	var targetMigrationID string
	if len(g.migrations) > 0 {
		targetMigrationID = g.migrations[len(g.migrations)-1].ID
	}
	return g.migrate(targetMigrationID)
}

// MigrateTo executes all migrations that did not run yet up to the migration that matches `migrationID`.
func (g *Gormigrate) MigrateTo(migrationID string) error {
	if err := g.checkIDExist(migrationID); err != nil {
		return err
	}
	return g.migrate(migrationID)
}

func (g *Gormigrate) migrate(migrationID string) error {
	if !g.hasMigrations() {
		return ErrNoMigrationDefined
	}

	if err := g.checkReservedID(); err != nil {
		return err
	}

	if err := g.checkDuplicatedID(); err != nil {
		return err
	}

	g.begin()
	defer g.rollback()

	if err := g.createMigrationTableIfNotExists(); err != nil {
		return err
	}

	if g.initSchema != nil && g.canInitializeSchema() {
		if err := g.runInitSchema(); err != nil {
			return err
		}
		return g.commit()
	}

	for _, migration := range g.migrations {
		if err := g.runMigration(migration); err != nil {
			return err
		}
		if migrationID != "" && migration.ID == migrationID {
			break
		}
	}
	return g.commit()
}

// There are migrations to apply if either there's a defined
// initSchema function or if the list of migrations is not empty.
func (g *Gormigrate) hasMigrations() bool {
	return g.initSchema != nil || len(g.migrations) > 0
}

// Check whether any migration is using a reserved ID.
// For now there's only have one reserved ID, but there may be more in the future.
func (g *Gormigrate) checkReservedID() error {
	for _, m := range g.migrations {
		if m.ID == initSchemaMigrationID {
			return &ReservedIDError{ID: m.ID}
		}
	}
	return nil
}

func (g *Gormigrate) checkDuplicatedID() error {
	lookup := make(map[string]struct{}, len(g.migrations))
	for _, m := range g.migrations {
		if _, ok := lookup[m.ID]; ok {
			return &DuplicatedIDError{ID: m.ID}
		}
		lookup[m.ID] = struct{}{}
	}
	return nil
}

func (g *Gormigrate) checkIDExist(migrationID string) error {
	for _, migrate := range g.migrations {
		if migrate.ID == migrationID {
			return nil
		}
	}
	return ErrMigrationIDDoesNotExist
}

// RollbackLast undo the last migration
func (g *Gormigrate) RollbackLast() error {
	if len(g.migrations) == 0 {
		return ErrNoMigrationDefined
	}

	g.begin()
	defer g.rollback()

	lastRunMigration, err := g.getLastRunMigration()
	if err != nil {
		return err
	}

	if err := g.rollbackMigration(lastRunMigration); err != nil {
		return err
	}
	return g.commit()
}

// RollbackTo undoes migrations up to the given migration that matches the `migrationID`.
// Migration with the matching `migrationID` is not rolled back.
func (g *Gormigrate) RollbackTo(migrationID string) error {
	if len(g.migrations) == 0 {
		return ErrNoMigrationDefined
	}

	if err := g.checkIDExist(migrationID); err != nil {
		return err
	}

	g.begin()
	defer g.rollback()

	for i := len(g.migrations) - 1; i >= 0; i-- {
		migration := g.migrations[i]
		if migration.ID == migrationID {
			break
		}
		if g.migrationDidRun(migration) {
			if err := g.rollbackMigration(migration); err != nil {
				return err
			}
		}
	}
	return g.commit()
}

func (g *Gormigrate) getLastRunMigration() (*Migration, error) {
	for i := len(g.migrations) - 1; i >= 0; i-- {
		migration := g.migrations[i]
		if g.migrationDidRun(migration) {
			return migration, nil
		}
	}
	return nil, ErrNoRunMigration
}

// RollbackMigration undo a migration.
func (g *Gormigrate) RollbackMigration(m *Migration) error {
	g.begin()
	defer g.rollback()

	if err := g.rollbackMigration(m); err != nil {
		return err
	}
	return g.commit()
}

func (g *Gormigrate) rollbackMigration(m *Migration) error {
	if m.Rollback == nil {
		return ErrRollbackImpossible
	}

	if err := m.Rollback(g.tx); err != nil {
		return err
	}

	sql := fmt.Sprintf("DELETE FROM %s WHERE %s = ?", g.options.TableName, g.options.IDColumnName)
	return g.tx.Exec(sql, m.ID).Error
}

func (g *Gormigrate) runInitSchema() error {
	if err := g.initSchema(g.tx); err != nil {
		return err
	}
	if err := g.insertMigration(initSchemaMigrationID); err != nil {
		return err
	}

	for _, migration := range g.migrations {
		if err := g.insertMigration(migration.ID); err != nil {
			return err
		}
	}

	return nil
}

func (g *Gormigrate) runMigration(migration *Migration) error {
	if len(migration.ID) == 0 {
		return ErrMissingID
	}

	if !g.migrationDidRun(migration) {
		if err := migration.Migrate(g.tx); err != nil {
			return err
		}

		if err := g.insertMigration(migration.ID); err != nil {
			return err
		}
	}
	return nil
}

func (g *Gormigrate) createMigrationTableIfNotExists() error {
	if g.tx.HasTable(g.options.TableName) {
		return nil
	}

	sql := fmt.Sprintf("CREATE TABLE %s (%s VARCHAR(%d) PRIMARY KEY)", g.options.TableName, g.options.IDColumnName, g.options.IDColumnSize)
	return g.tx.Exec(sql).Error
}

func (g *Gormigrate) migrationDidRun(m *Migration) bool {
	var count int
	g.tx.
		Table(g.options.TableName).
		Where(fmt.Sprintf("%s = ?", g.options.IDColumnName), m.ID).
		Count(&count)
	return count > 0
}

// The schema can be initialised only if it hasn't been initialised yet
// and no other migration has been applied already.
func (g *Gormigrate) canInitializeSchema() bool {
	if g.migrationDidRun(&Migration{ID: initSchemaMigrationID}) {
		return false
	}

	// If the ID doesn't exist, we also want the list of migrations to be empty
	var count int
	g.tx.
		Table(g.options.TableName).
		Count(&count)
	return count == 0
}

func (g *Gormigrate) insertMigration(id string) error {
	sql := fmt.Sprintf("INSERT INTO %s (%s) VALUES (?)", g.options.TableName, g.options.IDColumnName)
	return g.tx.Exec(sql, id).Error
}

func (g *Gormigrate) begin() {
	if g.options.UseTransaction {
		g.tx = g.db.Begin()
	} else {
		g.tx = g.db
	}
}

func (g *Gormigrate) commit() error {
	if g.options.UseTransaction {
		return g.tx.Commit().Error
	}
	return nil
}

func (g *Gormigrate) rollback() {
	if g.options.UseTransaction {
		g.tx.Rollback()
	}
}
<commit_msg>Check for errors<commit_after>package gormigrate

import (
	"errors"
	"fmt"

	"github.com/jinzhu/gorm"
)

const (
	initSchemaMigrationID = "SCHEMA_INIT"
)

// MigrateFunc is the func signature for migrating.
type MigrateFunc func(*gorm.DB) error

// RollbackFunc is the func signature for rollbacking.
type RollbackFunc func(*gorm.DB) error

// InitSchemaFunc is the func signature for initializing the schema.
type InitSchemaFunc func(*gorm.DB) error

// Options define options for all migrations.
type Options struct {
	// TableName is the migration table.
	TableName string
	// IDColumnName is the name of column where the migration id will be stored.
	IDColumnName string
	// IDColumnSize is the length of the migration id column
	IDColumnSize int
	// UseTransaction makes Gormigrate execute migrations inside a single transaction.
	// Keep in mind that not all databases support DDL commands inside transactions.
	UseTransaction bool
}

// Migration represents a database migration (a modification to be made on the database).
type Migration struct {
	// ID is the migration identifier. Usually a timestamp like "201601021504".
	ID string
	// Migrate is a function that will br executed while running this migration.
	Migrate MigrateFunc
	// Rollback will be executed on rollback. Can be nil.
	Rollback RollbackFunc
}

// Gormigrate represents a collection of all migrations of a database schema.
type Gormigrate struct {
	db         *gorm.DB
	tx         *gorm.DB
	options    *Options
	migrations []*Migration
	initSchema InitSchemaFunc
}

// ReservedIDError is returned when a migration is using a reserved ID
type ReservedIDError struct {
	ID string
}

func (e *ReservedIDError) Error() string {
	return fmt.Sprintf(`gormigrate: Reserved migration ID: "%s"`, e.ID)
}

// DuplicatedIDError is returned when more than one migration have the same ID
type DuplicatedIDError struct {
	ID string
}

func (e *DuplicatedIDError) Error() string {
	return fmt.Sprintf(`gormigrate: Duplicated migration ID: "%s"`, e.ID)
}

var (
	// DefaultOptions can be used if you don't want to think about options.
	DefaultOptions = &Options{
		TableName:      "migrations",
		IDColumnName:   "id",
		IDColumnSize:   255,
		UseTransaction: false,
	}

	// ErrRollbackImpossible is returned when trying to rollback a migration
	// that has no rollback function.
	ErrRollbackImpossible = errors.New("gormigrate: It's impossible to rollback this migration")

	// ErrNoMigrationDefined is returned when no migration is defined.
	ErrNoMigrationDefined = errors.New("gormigrate: No migration defined")

	// ErrMissingID is returned when the ID od migration is equal to ""
	ErrMissingID = errors.New("gormigrate: Missing ID in migration")

	// ErrNoRunMigration is returned when any run migration was found while
	// running RollbackLast
	ErrNoRunMigration = errors.New("gormigrate: Could not find last run migration")

	// ErrMigrationIDDoesNotExist is returned when migrating or rolling back to a migration ID that
	// does not exist in the list of migrations
	ErrMigrationIDDoesNotExist = errors.New("gormigrate: Tried to migrate to an ID that doesn't exist")
)

// New returns a new Gormigrate.
func New(db *gorm.DB, options *Options, migrations []*Migration) *Gormigrate {
	if options.TableName == "" {
		options.TableName = DefaultOptions.TableName
	}
	if options.IDColumnName == "" {
		options.IDColumnName = DefaultOptions.IDColumnName
	}
	if options.IDColumnSize == 0 {
		options.IDColumnSize = DefaultOptions.IDColumnSize
	}
	return &Gormigrate{
		db:         db,
		options:    options,
		migrations: migrations,
	}
}

// InitSchema sets a function that is run if no migration is found.
// The idea is preventing to run all migrations when a new clean database
// is being migrating. In this function you should create all tables and
// foreign key necessary to your application.
func (g *Gormigrate) InitSchema(initSchema InitSchemaFunc) {
	g.initSchema = initSchema
}

// Migrate executes all migrations that did not run yet.
func (g *Gormigrate) Migrate() error {
	if !g.hasMigrations() {
		return ErrNoMigrationDefined
	}
	var targetMigrationID string
	if len(g.migrations) > 0 {
		targetMigrationID = g.migrations[len(g.migrations)-1].ID
	}
	return g.migrate(targetMigrationID)
}

// MigrateTo executes all migrations that did not run yet up to the migration that matches `migrationID`.
func (g *Gormigrate) MigrateTo(migrationID string) error {
	if err := g.checkIDExist(migrationID); err != nil {
		return err
	}
	return g.migrate(migrationID)
}

func (g *Gormigrate) migrate(migrationID string) error {
	if !g.hasMigrations() {
		return ErrNoMigrationDefined
	}

	if err := g.checkReservedID(); err != nil {
		return err
	}

	if err := g.checkDuplicatedID(); err != nil {
		return err
	}

	g.begin()
	defer g.rollback()

	if err := g.createMigrationTableIfNotExists(); err != nil {
		return err
	}

	if g.initSchema != nil {
		canInitializeSchema, err := g.canInitializeSchema()
		if err != nil {
			return err
		}
		if canInitializeSchema {
			if err := g.runInitSchema(); err != nil {
				return err
			}
			return g.commit()
		}
	}

	for _, migration := range g.migrations {
		if err := g.runMigration(migration); err != nil {
			return err
		}
		if migrationID != "" && migration.ID == migrationID {
			break
		}
	}
	return g.commit()
}

// There are migrations to apply if either there's a defined
// initSchema function or if the list of migrations is not empty.
func (g *Gormigrate) hasMigrations() bool {
	return g.initSchema != nil || len(g.migrations) > 0
}

// Check whether any migration is using a reserved ID.
// For now there's only have one reserved ID, but there may be more in the future.
func (g *Gormigrate) checkReservedID() error {
	for _, m := range g.migrations {
		if m.ID == initSchemaMigrationID {
			return &ReservedIDError{ID: m.ID}
		}
	}
	return nil
}

func (g *Gormigrate) checkDuplicatedID() error {
	lookup := make(map[string]struct{}, len(g.migrations))
	for _, m := range g.migrations {
		if _, ok := lookup[m.ID]; ok {
			return &DuplicatedIDError{ID: m.ID}
		}
		lookup[m.ID] = struct{}{}
	}
	return nil
}

func (g *Gormigrate) checkIDExist(migrationID string) error {
	for _, migrate := range g.migrations {
		if migrate.ID == migrationID {
			return nil
		}
	}
	return ErrMigrationIDDoesNotExist
}

// RollbackLast undo the last migration
func (g *Gormigrate) RollbackLast() error {
	if len(g.migrations) == 0 {
		return ErrNoMigrationDefined
	}

	g.begin()
	defer g.rollback()

	lastRunMigration, err := g.getLastRunMigration()
	if err != nil {
		return err
	}

	if err := g.rollbackMigration(lastRunMigration); err != nil {
		return err
	}
	return g.commit()
}

// RollbackTo undoes migrations up to the given migration that matches the `migrationID`.
// Migration with the matching `migrationID` is not rolled back.
func (g *Gormigrate) RollbackTo(migrationID string) error {
	if len(g.migrations) == 0 {
		return ErrNoMigrationDefined
	}

	if err := g.checkIDExist(migrationID); err != nil {
		return err
	}

	g.begin()
	defer g.rollback()

	for i := len(g.migrations) - 1; i >= 0; i-- {
		migration := g.migrations[i]
		if migration.ID == migrationID {
			break
		}
		migrationRan, err := g.migrationRan(migration)
		if err != nil {
			return err
		}
		if migrationRan {
			if err := g.rollbackMigration(migration); err != nil {
				return err
			}
		}
	}
	return g.commit()
}

func (g *Gormigrate) getLastRunMigration() (*Migration, error) {
	for i := len(g.migrations) - 1; i >= 0; i-- {
		migration := g.migrations[i]

		migrationRan, err := g.migrationRan(migration)
		if err != nil {
			return nil, err
		}

		if migrationRan {
			return migration, nil
		}
	}
	return nil, ErrNoRunMigration
}

// RollbackMigration undo a migration.
func (g *Gormigrate) RollbackMigration(m *Migration) error {
	g.begin()
	defer g.rollback()

	if err := g.rollbackMigration(m); err != nil {
		return err
	}
	return g.commit()
}

func (g *Gormigrate) rollbackMigration(m *Migration) error {
	if m.Rollback == nil {
		return ErrRollbackImpossible
	}

	if err := m.Rollback(g.tx); err != nil {
		return err
	}

	sql := fmt.Sprintf("DELETE FROM %s WHERE %s = ?", g.options.TableName, g.options.IDColumnName)
	return g.tx.Exec(sql, m.ID).Error
}

func (g *Gormigrate) runInitSchema() error {
	if err := g.initSchema(g.tx); err != nil {
		return err
	}
	if err := g.insertMigration(initSchemaMigrationID); err != nil {
		return err
	}

	for _, migration := range g.migrations {
		if err := g.insertMigration(migration.ID); err != nil {
			return err
		}
	}

	return nil
}

func (g *Gormigrate) runMigration(migration *Migration) error {
	if len(migration.ID) == 0 {
		return ErrMissingID
	}

	migrationRan, err := g.migrationRan(migration)
	if err != nil {
		return err
	}
	if !migrationRan {
		if err := migration.Migrate(g.tx); err != nil {
			return err
		}

		if err := g.insertMigration(migration.ID); err != nil {
			return err
		}
	}
	return nil
}

func (g *Gormigrate) createMigrationTableIfNotExists() error {
	if g.tx.HasTable(g.options.TableName) {
		return nil
	}

	sql := fmt.Sprintf("CREATE TABLE %s (%s VARCHAR(%d) PRIMARY KEY)", g.options.TableName, g.options.IDColumnName, g.options.IDColumnSize)
	return g.tx.Exec(sql).Error
}

func (g *Gormigrate) migrationRan(m *Migration) (bool, error) {
	var count int
	err := g.tx.
		Table(g.options.TableName).
		Where(fmt.Sprintf("%s = ?", g.options.IDColumnName), m.ID).
		Count(&count).
		Error
	return count > 0, err
}

// The schema can be initialised only if it hasn't been initialised yet
// and no other migration has been applied already.
func (g *Gormigrate) canInitializeSchema() (bool, error) {
	migrationRan, err := g.migrationRan(&Migration{ID: initSchemaMigrationID})
	if err != nil {
		return false, err
	}
	if migrationRan {
		return false, nil
	}

	// If the ID doesn't exist, we also want the list of migrations to be empty
	var count int
	err = g.tx.
		Table(g.options.TableName).
		Count(&count).
		Error
	return count == 0, err
}

func (g *Gormigrate) insertMigration(id string) error {
	sql := fmt.Sprintf("INSERT INTO %s (%s) VALUES (?)", g.options.TableName, g.options.IDColumnName)
	return g.tx.Exec(sql, id).Error
}

func (g *Gormigrate) begin() {
	if g.options.UseTransaction {
		g.tx = g.db.Begin()
	} else {
		g.tx = g.db
	}
}

func (g *Gormigrate) commit() error {
	if g.options.UseTransaction {
		return g.tx.Commit().Error
	}
	return nil
}

func (g *Gormigrate) rollback() {
	if g.options.UseTransaction {
		g.tx.Rollback()
	}
}
<|endoftext|>


<commit_before>package prompt

import (
	"strconv"
	"syscall"
)

type VT100Writer struct {
	fd     int
	buffer []byte
}

func (w *VT100Writer) Write(data []byte) {
	w.WriteRaw(byteFilter(data, writeFilter))
	return
}

func (w *VT100Writer) WriteStr(data string) {
	w.Write([]byte(data))
	return
}

func (w *VT100Writer) WriteRaw(data []byte) {
	w.buffer = append(w.buffer, data...)
	return
}

func (w *VT100Writer) Flush() error {
	_, err := syscall.Write(w.fd, w.buffer)
	if err != nil {
		return err
	}
	w.buffer = []byte{}
	return nil
}

func (w *VT100Writer) Clear() {
	syscall.Write(w.fd, []byte{0x1b, 0x5b, 0x02, 0x6a, 0x1b, 0x63})
	return
}

/* Cursor */

func (w *VT100Writer) CursorUp(n int) {
	if n < 0 {
		w.CursorDown(n)
		return
	}
	s := strconv.Itoa(n)
	w.WriteRaw([]byte{0x1b, 0x5b})
	w.WriteRaw([]byte(s))
	w.WriteRaw([]byte{0x41})
	return
}

func (w *VT100Writer) CursorDown(n int) {
	if n < 0 {
		w.CursorUp(n)
		return
	}
	s := strconv.Itoa(n)
	w.WriteRaw([]byte{0x1b, 0x5b})
	w.WriteRaw([]byte(s))
	w.WriteRaw([]byte{0x42})
	return
}

func (w *VT100Writer) CursorForward(n int) {
	if n < 0 {
		w.CursorBackward(n)
		return
	}
	s := strconv.Itoa(n)
	w.WriteRaw([]byte{0x1b, 0x5b})
	w.WriteRaw([]byte(s))
	w.WriteRaw([]byte{0x43})
	return
}

func (w *VT100Writer) CursorBackward(n int) {
	if n < 0 {
		w.CursorForward(n)
		return
	}
	s := strconv.Itoa(n)
	w.WriteRaw([]byte{0x1b, 0x5b})
	w.WriteRaw([]byte(s))
	w.WriteRaw([]byte{0x44})
	return
}

/* Title */

func (w *VT100Writer) SetTitle(title string) {
	w.WriteRaw([]byte{0x1b, 0x5d, 0x32, 0x3b})
	w.WriteRaw(byteFilter([]byte(title), setTextFilter))
	w.WriteRaw([]byte{0x07})
	return
}

func (w *VT100Writer) ClearTitle() {
	w.WriteRaw([]byte{0x1b, 0x5d, 0x32, 0x3b, 0x07})
	return
}

func NewVT100Writer() *VT100Writer {
	return &VT100Writer{
		fd: syscall.Stdout,
	}
}

func writeFilter(buf byte) bool {
	return buf != 0x1b && buf != 0x3f
}

func setTextFilter(buf byte) bool {
	return buf != 0x1b && buf != 0x07
}

func byteFilter(buf []byte, fn ...func(b byte) bool) []byte {
	if len(fn) == 0 {
		return buf
	}
	ret := make([]byte, 0, len(buf))
	f := fn[0]
	for i, n := range buf {
		if f(n) {
			ret = append(ret, buf[i])
		}
	}
	return byteFilter(ret, fn[1:]...)
}
<commit_msg>Add erase operation in vt100 output<commit_after>package prompt

import (
	"strconv"
	"syscall"
)

type VT100Writer struct {
	fd     int
	buffer []byte
}

func (w *VT100Writer) Write(data []byte) {
	w.WriteRaw(byteFilter(data, writeFilter))
	return
}

func (w *VT100Writer) WriteStr(data string) {
	w.Write([]byte(data))
	return
}

func (w *VT100Writer) WriteRaw(data []byte) {
	w.buffer = append(w.buffer, data...)
	return
}

func (w *VT100Writer) Flush() error {
	_, err := syscall.Write(w.fd, w.buffer)
	if err != nil {
		return err
	}
	w.buffer = []byte{}
	return nil
}

/* Erase */

func (w *VT100Writer) Clear() {
	w.WriteRaw([]byte{0x1b, 0x5b, 0x02, 0x6a, 0x1b, 0x63})
	return
}

func (w *VT100Writer) EraseScreen() {
	w.WriteRaw([]byte{0x1b, 0x5b, 0x32, 0x4a})
	return
}

func (w *VT100Writer) EraseDown() {
	w.WriteRaw([]byte{0x1b, 0x5b, 0x4a})
	return
}

func (w *VT100Writer) EraseEndOfLine() {
	w.WriteRaw([]byte{0x1b, 0x5b, 0x4b})
	return
}

/* Cursor */

func (w *VT100Writer) ShowCursor() {
	w.WriteRaw([]byte{0x1b, 0x5b, 0x3f, 0x31, 0x32, 0x6c, 0x1b, 0x5b, 0x3f, 0x32, 0x35, 0x68})
}

func (w *VT100Writer) HideCursor() {
	w.WriteRaw([]byte{0x1b, 0x5b, 0x3f, 0x32, 0x35, 0x6c})
	return
}

func (w *VT100Writer) CursorGoTo(row, col int) {
	r := strconv.Itoa(row)
	c := strconv.Itoa(col)
	w.WriteRaw([]byte{0x1b, 0x5b})
	w.WriteRaw([]byte(r))
	w.WriteRaw([]byte{0x3b})
	w.WriteRaw([]byte(c))
	w.WriteRaw([]byte{0x48})
	return
}

func (w *VT100Writer) CursorUp(n int) {
	if n < 0 {
		w.CursorDown(n)
		return
	}
	s := strconv.Itoa(n)
	w.WriteRaw([]byte{0x1b, 0x5b})
	w.WriteRaw([]byte(s))
	w.WriteRaw([]byte{0x41})
	return
}

func (w *VT100Writer) CursorDown(n int) {
	if n < 0 {
		w.CursorUp(n)
		return
	}
	s := strconv.Itoa(n)
	w.WriteRaw([]byte{0x1b, 0x5b})
	w.WriteRaw([]byte(s))
	w.WriteRaw([]byte{0x42})
	return
}

func (w *VT100Writer) CursorForward(n int) {
	if n < 0 {
		w.CursorBackward(n)
		return
	}
	s := strconv.Itoa(n)
	w.WriteRaw([]byte{0x1b, 0x5b})
	w.WriteRaw([]byte(s))
	w.WriteRaw([]byte{0x43})
	return
}

func (w *VT100Writer) CursorBackward(n int) {
	if n < 0 {
		w.CursorForward(n)
		return
	}
	s := strconv.Itoa(n)
	w.WriteRaw([]byte{0x1b, 0x5b})
	w.WriteRaw([]byte(s))
	w.WriteRaw([]byte{0x44})
	return
}

/* Title */

func (w *VT100Writer) SetTitle(title string) {
	w.WriteRaw([]byte{0x1b, 0x5d, 0x32, 0x3b})
	w.WriteRaw(byteFilter([]byte(title), setTextFilter))
	w.WriteRaw([]byte{0x07})
	return
}

func (w *VT100Writer) ClearTitle() {
	w.WriteRaw([]byte{0x1b, 0x5d, 0x32, 0x3b, 0x07})
	return
}

func NewVT100Writer() *VT100Writer {
	return &VT100Writer{
		fd: syscall.Stdout,
	}
}

func writeFilter(buf byte) bool {
	return buf != 0x1b && buf != 0x3f
}

func setTextFilter(buf byte) bool {
	return buf != 0x1b && buf != 0x07
}

func byteFilter(buf []byte, fn ...func(b byte) bool) []byte {
	if len(fn) == 0 {
		return buf
	}
	ret := make([]byte, 0, len(buf))
	f := fn[0]
	for i, n := range buf {
		if f(n) {
			ret = append(ret, buf[i])
		}
	}
	return byteFilter(ret, fn[1:]...)
}
<|endoftext|>


<commit_before>package redis_timeseries

import (
	"time"

	"github.com/jinzhu/now"
)

// Get returns unix timestamps of the times closest to each time argument
// as specified by the interval.
//
// Ex:
//   newyear := time.Date(2000, time.January, 01, 0, 0, 30, 0, time.UTC)
//   newyearplus := time.Date(2000, time.January, 01, 0, 16, 0, 0, time.UTC)
//
//   results := Get(15*time.Minute, newyear, newyearplus)
//
// `newyear` falls between 00:00 and 00:15, so it returns 00:00.
// `newyearplus` falls between 00:15 and 00:30, so it returns 00:15.
//
//  If a time falls on the interval, ie 00:15, it returns 00.15;
func Get(interval time.Duration, times ...time.Time) []int64 {
	var results = []int64{}

	for _, t := range times {
		start := now.New(t).BeginningOfHour()

		for {
			// if t falls on start+interval, then we add start+interval
			if t.Equal(start.Add(interval)) {
				results = append(results, start.Add(interval).Unix())
				break
			}

			// if t isn't greater than start+interval, then we add start
			if !t.After(start.Add(interval)) {
				results = append(results, start.Unix())
				break
			}

			// this means t must be greater than start+interval, however
			// we don't know by how much, so we continue with the loop
			start = start.Add(interval)
		}
	}

	return results
}
<commit_msg>comments format fix<commit_after>package redis_timeseries

import (
	"time"

	"github.com/jinzhu/now"
)

// Get returns unix timestamps of the times closest to each time argument
// as specified by the interval.
//
// Ex:
//   newyear := time.Date(2000, time.January, 01, 0, 0, 30, 0, time.UTC)
//   newyearplus := time.Date(2000, time.January, 01, 0, 16, 0, 0, time.UTC)
//
//   // returns: []string{946684800, 946685700}
//   results := Get(15*time.Minute, newyear, newyearplus)
//
// `newyear` falls between 00:00 and 00:15, so it returns 00:00.
// `newyearplus` falls between 00:15 and 00:30, so it returns 00:15.
//
// If a time falls on the interval, ie 00:15, it returns 00.15.
func Get(interval time.Duration, times ...time.Time) []int64 {
	var results = []int64{}

	for _, t := range times {
		start := now.New(t).BeginningOfHour()

		for {
			// if t falls on start+interval, then we add start+interval
			if t.Equal(start.Add(interval)) {
				results = append(results, start.Add(interval).Unix())
				break
			}

			// if t isn't greater than start+interval, then we add start
			if !t.After(start.Add(interval)) {
				results = append(results, start.Unix())
				break
			}

			// this means t must be greater than start+interval, however
			// we don't know by how much, so we continue with the loop
			start = start.Add(interval)
		}
	}

	return results
}
<|endoftext|>


<commit_before>package contracthash

import (
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha256"
	"errors"
	"fmt"
	"io"
	"math/big"

	"github.com/conformal/btcec"
	"github.com/conformal/btcnet"
	"github.com/conformal/btcscript"
	"github.com/conformal/btcutil"
	"github.com/conformal/btcwire"
)

// secp256k1
var curve = btcec.S256()
var secp256k1_Order = curve.Params().N

// DeriveClaimKey derives a private key based on a contract hash and a
// functionary's public key, allowing a member of the federation to sign-off on
// redemption of the output payed to the contract hash.
func DeriveClaimKey(privKey *btcutil.WIF, nonce []byte, addr *btcutil.Address, netParams *btcnet.Params) (*btcutil.WIF, error) {
	contractData, err := computeContractData(nonce, *addr)
	if err != nil {
		return nil, err
	}

	tweak, err := computeTweak(privKey.PrivKey.PubKey(), contractData)
	if err != nil {
		return nil, err
	}

	tweakInt := new(big.Int).SetBytes(tweak.Bytes())

	// derivedKey = tweak + privKey
	derivedSecret := new(big.Int).Mod(
		new(big.Int).Add(tweakInt, privKey.PrivKey.D), secp256k1_Order)
	claimKey, _ := btcec.PrivKeyFromBytes(curve, derivedSecret.Bytes())

	wifKey, err := btcutil.NewWIF(claimKey, netParams, true)
	if err != nil {
		return nil, err
	}

	return wifKey, err
}

// DeriveContractHash homomorphically derives a pay-to-contract-address
// based on the passed nonce and contractTemplate (a multi-sig
// redeemScript with the federations public keys). The derivation method
// is described in Appendix A of the sidechains whitepaper.
func DeriveContractHash(contractTemplate []byte, addr *btcutil.Address, nonce *[]byte, netParams *btcnet.Params) (*Contract, error) {
	// Extract info about the contract template redeemscript.
	scriptType, addrs, numRequiredSigs, err := btcscript.ExtractPkScriptAddrs(
		contractTemplate, &btcnet.TestNet3Params)
	if err != nil {
		return nil, err
	}

	// Abort if the redeemscript isn't multi-sig.
	if scriptType != btcscript.MultiSigTy {
		return nil, errors.New(
			"The contractTemplate must be a multi-sig redeemscript")
	}

	// Read a fresh 16-byte nonce from the OS if one isn't provided.
	if nonce == nil {
		nonce := make([]byte, 0, 16)
		_, err = io.ReadFull(rand.Reader, nonce)
		if err != nil {
			return nil, errors.New(fmt.Sprintf("Error reading random nonce: ", err))
		}
	}

	// The contract whose hash we'd like to pay to.
	contractData, err := computeContractData(*nonce, *addr)
	if err != nil {
		return nil, err
	}

	// Extract pubkeys from the wrapper type.
	fedPubKeys := make([]*btcec.PublicKey, 0, len(addrs))
	for _, addr := range addrs {
		fedPubKeys = append(fedPubKeys, addr.(*btcutil.AddressPubKey).PubKey())
	}

	derivedKeys := make([]*btcutil.AddressPubKey, 0, len(fedPubKeys))
	for _, fedKey := range fedPubKeys {
		tweak, err := computeTweak(fedKey, contractData)
		if err != nil {
			return nil, err
		}

		tweakInt := new(big.Int).SetBytes(tweak.Bytes())
		if tweakInt.Cmp(secp256k1_Order) > 0 {
			return nil, errors.New("Current tweak exceeds order, " +
				"pick a new nonce!")
		}

		// PCC_i = P_i + G x Tweak_i
		tX, tY := curve.ScalarBaseMult(tweak.Bytes())
		nX, nY := curve.Add(tX, tY, fedKey.X, fedKey.Y)
		point := (&btcec.PublicKey{curve, nX, nY}).SerializeCompressed()
		derivedKey, _ := btcutil.NewAddressPubKey(point, netParams)
		derivedKeys = append(derivedKeys, derivedKey)
	}

	script, err := btcscript.MultiSigScript(derivedKeys, numRequiredSigs)
	if err != nil {
		return nil, err
	}

	p2shAddr, err := btcutil.NewAddressScriptHash(script, netParams)
	if err != nil {
		return nil, err
	}

	contract := &Contract{
		Nonce:             *nonce,
		RedeemScript:      script,
		PayToContractAddr: p2shAddr,
	}

	return contract, nil
}

type Contract struct {
	Nonce             []byte
	RedeemScript      []byte
	PayToContractAddr *btcutil.AddressScriptHash
}

// computeTweak computes a unique tweak for a given federation key and contract.
// The tweak is the output of a HMAC-SHA256 with the fed pub keys as our secret,
// and the contract as our message.
func computeTweak(fedKey *btcec.PublicKey, contract []byte) (*btcwire.ShaHash, error) {
	mac := hmac.New(sha256.New, fedKey.SerializeCompressed())
	mac.Write(contract)
	digest := mac.Sum(nil)

	sha, err := btcwire.NewShaHash(digest)
	if err != nil {
		return nil, err
	}

	return sha, nil
}

// computeContractData creates the contract who's hash will be payed to.
// Our contract is 40 bytes of data: four bytes of a four character abreviation
// of passed address type, 16 bytes for the passed nonce, and finally 20 bytes
// for the decoded base58 address.
func computeContractData(nonce []byte, addr btcutil.Address) ([]byte, error) {
	var addressType string
	switch addr.(type) {
	case *btcutil.AddressScriptHash:
		addressType = "P2SH"
	case *btcutil.AddressPubKeyHash:
		addressType = "P2PH"
	default:
		addressType = "UNKNOWN"
	}

	if addressType == "UNKNOWN" {
		return nil, errors.New("Unsupported address type")
	}

	hashBytes := addr.ScriptAddress()

	contractData := make([]byte, 0, 40)
	contractData = append(contractData, []byte(addressType)...)
	contractData = append(contractData, nonce...)
	contractData = append(contractData, hashBytes...)

	return contractData, nil
}
<commit_msg>go vet<commit_after>package contracthash

import (
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha256"
	"errors"
	"fmt"
	"io"
	"math/big"

	"github.com/conformal/btcec"
	"github.com/conformal/btcnet"
	"github.com/conformal/btcscript"
	"github.com/conformal/btcutil"
	"github.com/conformal/btcwire"
)

// secp256k1
var curve = btcec.S256()
var secp256k1_Order = curve.Params().N

// DeriveClaimKey derives a private key based on a contract hash and a
// functionary's public key, allowing a member of the federation to sign-off on
// redemption of the output payed to the contract hash.
func DeriveClaimKey(privKey *btcutil.WIF, nonce []byte, addr *btcutil.Address, netParams *btcnet.Params) (*btcutil.WIF, error) {
	contractData, err := computeContractData(nonce, *addr)
	if err != nil {
		return nil, err
	}

	tweak, err := computeTweak(privKey.PrivKey.PubKey(), contractData)
	if err != nil {
		return nil, err
	}

	tweakInt := new(big.Int).SetBytes(tweak.Bytes())

	// derivedKey = tweak + privKey
	derivedSecret := new(big.Int).Mod(
		new(big.Int).Add(tweakInt, privKey.PrivKey.D), secp256k1_Order)
	claimKey, _ := btcec.PrivKeyFromBytes(curve, derivedSecret.Bytes())

	wifKey, err := btcutil.NewWIF(claimKey, netParams, true)
	if err != nil {
		return nil, err
	}

	return wifKey, err
}

// DeriveContractHash homomorphically derives a pay-to-contract-address
// based on the passed nonce and contractTemplate (a multi-sig
// redeemScript with the federations public keys). The derivation method
// is described in Appendix A of the sidechains whitepaper.
func DeriveContractHash(contractTemplate []byte, addr *btcutil.Address, nonce *[]byte, netParams *btcnet.Params) (*Contract, error) {
	// Extract info about the contract template redeemscript.
	scriptType, addrs, numRequiredSigs, err := btcscript.ExtractPkScriptAddrs(
		contractTemplate, &btcnet.TestNet3Params)
	if err != nil {
		return nil, err
	}

	// Abort if the redeemscript isn't multi-sig.
	if scriptType != btcscript.MultiSigTy {
		return nil, errors.New(
			"The contractTemplate must be a multi-sig redeemscript")
	}

	// Read a fresh 16-byte nonce from the OS if one isn't provided.
	if nonce == nil {
		nonce := make([]byte, 0, 16)
		_, err = io.ReadFull(rand.Reader, nonce)
		if err != nil {
			return nil, errors.New(fmt.Sprintf("Error reading random nonce: %v", err))
		}
	}

	// The contract whose hash we'd like to pay to.
	contractData, err := computeContractData(*nonce, *addr)
	if err != nil {
		return nil, err
	}

	// Extract pubkeys from the wrapper type.
	fedPubKeys := make([]*btcec.PublicKey, 0, len(addrs))
	for _, addr := range addrs {
		fedPubKeys = append(fedPubKeys, addr.(*btcutil.AddressPubKey).PubKey())
	}

	derivedKeys := make([]*btcutil.AddressPubKey, 0, len(fedPubKeys))
	for _, fedKey := range fedPubKeys {
		tweak, err := computeTweak(fedKey, contractData)
		if err != nil {
			return nil, err
		}

		tweakInt := new(big.Int).SetBytes(tweak.Bytes())
		if tweakInt.Cmp(secp256k1_Order) > 0 {
			return nil, errors.New("Current tweak exceeds order, " +
				"pick a new nonce!")
		}

		// PCC_i = P_i + G x Tweak_i
		tX, tY := curve.ScalarBaseMult(tweak.Bytes())
		nX, nY := curve.Add(tX, tY, fedKey.X, fedKey.Y)
		point := (&btcec.PublicKey{Curve: curve, X: nX, Y: nY}).SerializeCompressed()
		derivedKey, _ := btcutil.NewAddressPubKey(point, netParams)
		derivedKeys = append(derivedKeys, derivedKey)
	}

	script, err := btcscript.MultiSigScript(derivedKeys, numRequiredSigs)
	if err != nil {
		return nil, err
	}

	p2shAddr, err := btcutil.NewAddressScriptHash(script, netParams)
	if err != nil {
		return nil, err
	}

	contract := &Contract{
		Nonce:             *nonce,
		RedeemScript:      script,
		PayToContractAddr: p2shAddr,
	}

	return contract, nil
}

type Contract struct {
	Nonce             []byte
	RedeemScript      []byte
	PayToContractAddr *btcutil.AddressScriptHash
}

// computeTweak computes a unique tweak for a given federation key and contract.
// The tweak is the output of a HMAC-SHA256 with the fed pub keys as our secret,
// and the contract as our message.
func computeTweak(fedKey *btcec.PublicKey, contract []byte) (*btcwire.ShaHash, error) {
	mac := hmac.New(sha256.New, fedKey.SerializeCompressed())
	mac.Write(contract)
	digest := mac.Sum(nil)

	sha, err := btcwire.NewShaHash(digest)
	if err != nil {
		return nil, err
	}

	return sha, nil
}

// computeContractData creates the contract who's hash will be payed to.
// Our contract is 40 bytes of data: four bytes of a four character abreviation
// of passed address type, 16 bytes for the passed nonce, and finally 20 bytes
// for the decoded base58 address.
func computeContractData(nonce []byte, addr btcutil.Address) ([]byte, error) {
	var addressType string
	switch addr.(type) {
	case *btcutil.AddressScriptHash:
		addressType = "P2SH"
	case *btcutil.AddressPubKeyHash:
		addressType = "P2PH"
	default:
		addressType = "UNKNOWN"
	}

	if addressType == "UNKNOWN" {
		return nil, errors.New("Unsupported address type")
	}

	hashBytes := addr.ScriptAddress()

	contractData := make([]byte, 0, 40)
	contractData = append(contractData, []byte(addressType)...)
	contractData = append(contractData, nonce...)
	contractData = append(contractData, hashBytes...)

	return contractData, nil
}
<|endoftext|>


<commit_before>// Copyright 2013 tsuru authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package provision provides interfaces that need to be satisfied in order to
// implement a new provisioner on tsuru.
package provision

import (
	"fmt"
	"io"
)

type Status string

func (s Status) String() string {
	return string(s)
}

const (
	StatusStarted    = Status("started")
	StatusPending    = Status("pending")
	StatusDown       = Status("down")
	StatusError      = Status("error")
	StatusInstalling = Status("installing")
	StatusCreating   = Status("creating")
)

// Unit represents a provision unit. Can be a machine, container or anything
// IP-addressable.
type Unit struct {
	Name       string
	AppName    string
	Type       string
	InstanceId string
	Machine    int
	Ip         string
	Status     Status
}

// AppUnit represents a unit in an app.
type AppUnit interface {
	// Returns the name of the unit.
	GetName() string

	// Returns the number of the unit.
	GetMachine() int

	// Returns the status of the unit.
	GetStatus() Status
}

// App represents a tsuru app.
//
// It contains only relevant information for provisioning.
type App interface {
	// Log should be used to log messages in the app.
	Log(message, source string) error

	// GetName returns the name of the app.
	GetName() string

	// GetFramework returns the framework (type) of the app. It is equivalent
	// to the Unit `Type` field.
	GetFramework() string

	// GetUnits returns all units of the app, in a slice.
	ProvisionUnits() []AppUnit
}

// Provisioner is the basic interface of this package.
//
// Any tsuru provisioner must implement this interface in order to provision
// tsuru apps.
//
// Tsuru comes with a default provisioner: juju. One can add other provisioners
// by satisfying this interface and registering it using the function Register.
type Provisioner interface {
	// Provision is called when tsuru is creating the app.
	Provision(App) error

	// Destroy is called when tsuru is destroying the app.
	Destroy(App) error

	// AddUnits adds units to an app. The first parameter is the app, the
	// second is the number of units to add.
	//
	// It returns a slice containing all added units
	AddUnits(App, uint) ([]Unit, error)

	// RemoveUnit removes a unit from the app. It receives the app and the name
	// of the unit to be removed.
	RemoveUnit(App, string) error

	// RemoveUnits removes multiple units from an app. The first parameter it
	// the app, the second is the number of units to remove.
	//
	// It returns a slice containing indices of all removed units (the index
	// must match the slice returned by App.ProvisionUnits). The list of
	// indices must be returned sorted.
	RemoveUnits(App, uint) ([]int, error)

	// ExecuteCommand runs a command in all units of the app.
	ExecuteCommand(stdout, stderr io.Writer, app App, cmd string, args ...string) error

	// CollectStatus returns information about all provisioned units. It's used
	// by tsuru collector when updating the status of apps in the database.
	CollectStatus() ([]Unit, error)

	// Returns a load balancer manager.
	LoadBalancer() LBManager
}

// LBManager provides methods to manage load balancers per app.
type LBManager interface {
	// Creates a new load balancer.
	Create(App) error

	// Destroys a load balancer.
	Destroy(App) error

	// Adds a unit to the load balancer.
	Register(App, Unit) error

	// Removes a unit from the load balancer.
	Deregister(App, Unit) error

	// Addr returns the address of the load balancer for an app. It will
	// probably be a DNS name or IP address.
	Addr(App) (string, error)
}

var provisioners = make(map[string]Provisioner)

// Register registers a new provisioner in the Provisioner registry.
func Register(name string, p Provisioner) {
	provisioners[name] = p
}

// Get gets the named provisioner from the registry.
func Get(name string) (Provisioner, error) {
	p, ok := provisioners[name]
	if !ok {
		return nil, fmt.Errorf("Unknown provisioner: %q.", name)
	}
	return p, nil
}

type Error struct {
	Reason string
	Err    error
}

func (e *Error) Error() string {
	var err string
	if e.Err != nil {
		err = e.Err.Error() + ": " + e.Reason
	} else {
		err = e.Reason
	}
	return err
}
<commit_msg>provision: document that Addr will be used by collector<commit_after>// Copyright 2013 tsuru authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package provision provides interfaces that need to be satisfied in order to
// implement a new provisioner on tsuru.
package provision

import (
	"fmt"
	"io"
)

type Status string

func (s Status) String() string {
	return string(s)
}

const (
	StatusStarted    = Status("started")
	StatusPending    = Status("pending")
	StatusDown       = Status("down")
	StatusError      = Status("error")
	StatusInstalling = Status("installing")
	StatusCreating   = Status("creating")
)

// Unit represents a provision unit. Can be a machine, container or anything
// IP-addressable.
type Unit struct {
	Name       string
	AppName    string
	Type       string
	InstanceId string
	Machine    int
	Ip         string
	Status     Status
}

// AppUnit represents a unit in an app.
type AppUnit interface {
	// Returns the name of the unit.
	GetName() string

	// Returns the number of the unit.
	GetMachine() int

	// Returns the status of the unit.
	GetStatus() Status
}

// App represents a tsuru app.
//
// It contains only relevant information for provisioning.
type App interface {
	// Log should be used to log messages in the app.
	Log(message, source string) error

	// GetName returns the name of the app.
	GetName() string

	// GetFramework returns the framework (type) of the app. It is equivalent
	// to the Unit `Type` field.
	GetFramework() string

	// GetUnits returns all units of the app, in a slice.
	ProvisionUnits() []AppUnit
}

// Provisioner is the basic interface of this package.
//
// Any tsuru provisioner must implement this interface in order to provision
// tsuru apps.
//
// Tsuru comes with a default provisioner: juju. One can add other provisioners
// by satisfying this interface and registering it using the function Register.
type Provisioner interface {
	// Provision is called when tsuru is creating the app.
	Provision(App) error

	// Destroy is called when tsuru is destroying the app.
	Destroy(App) error

	// AddUnits adds units to an app. The first parameter is the app, the
	// second is the number of units to add.
	//
	// It returns a slice containing all added units
	AddUnits(App, uint) ([]Unit, error)

	// RemoveUnit removes a unit from the app. It receives the app and the name
	// of the unit to be removed.
	RemoveUnit(App, string) error

	// RemoveUnits removes multiple units from an app. The first parameter it
	// the app, the second is the number of units to remove.
	//
	// It returns a slice containing indices of all removed units (the index
	// must match the slice returned by App.ProvisionUnits). The list of
	// indices must be returned sorted.
	RemoveUnits(App, uint) ([]int, error)

	// ExecuteCommand runs a command in all units of the app.
	ExecuteCommand(stdout, stderr io.Writer, app App, cmd string, args ...string) error

	// CollectStatus returns information about all provisioned units. It's used
	// by tsuru collector when updating the status of apps in the database.
	CollectStatus() ([]Unit, error)

	// Returns a load balancer manager.
	LoadBalancer() LBManager
}

// LBManager provides methods to manage load balancers per app.
type LBManager interface {
	// Creates a new load balancer.
	Create(App) error

	// Destroys a load balancer.
	Destroy(App) error

	// Adds a unit to the load balancer.
	Register(App, Unit) error

	// Removes a unit from the load balancer.
	Deregister(App, Unit) error

	// Addr returns the address of the load balancer for an app. It will
	// probably be a DNS name or IP address.
	//
	// Collector will use this method to get the IP (althought it might not
	// be an actual IP, collector calls it "IP") of the app from the
	// provisioner.
	Addr(App) (string, error)
}

var provisioners = make(map[string]Provisioner)

// Register registers a new provisioner in the Provisioner registry.
func Register(name string, p Provisioner) {
	provisioners[name] = p
}

// Get gets the named provisioner from the registry.
func Get(name string) (Provisioner, error) {
	p, ok := provisioners[name]
	if !ok {
		return nil, fmt.Errorf("Unknown provisioner: %q.", name)
	}
	return p, nil
}

type Error struct {
	Reason string
	Err    error
}

func (e *Error) Error() string {
	var err string
	if e.Err != nil {
		err = e.Err.Error() + ": " + e.Reason
	} else {
		err = e.Reason
	}
	return err
}
<|endoftext|>


<commit_before>// Copyright © 2016 The Things Network
// Use of this source code is governed by the MIT license that can be found in the LICENSE file.

package handler

import (
	"bytes"
	"encoding/binary"
	"reflect"
	"time"

	. "github.com/TheThingsNetwork/ttn/core"
	"github.com/TheThingsNetwork/ttn/utils/errors"
	"github.com/TheThingsNetwork/ttn/utils/pointer"
	"github.com/TheThingsNetwork/ttn/utils/stats"
	"github.com/apex/log"
	"github.com/brocaar/lorawan"
)

const bufferDelay time.Duration = time.Millisecond * 300

// component implements the core.Component interface
type component struct {
	broker  JSONRecipient
	ctx     log.Interface
	devices DevStorage
	packets PktStorage
	set     chan<- bundle
}

type bundle struct {
	Adapter Adapter
	Chresp  chan interface{}
	Entry   devEntry
	ID      [20]byte
	Packet  HPacket
	Time    time.Time
}

// New construct a new Handler
func New(devDb DevStorage, pktDb PktStorage, broker JSONRecipient, ctx log.Interface) Handler {
	h := component{
		ctx:     ctx,
		devices: devDb,
		packets: pktDb,
		broker:  broker,
	}

	set := make(chan bundle)
	bundles := make(chan []bundle)

	h.set = set
	go h.consumeBundles(bundles)
	go h.consumeSet(bundles, set)

	return h
}

// Register implements the core.Component interface
func (h component) Register(reg Registration, an AckNacker, sub Subscriber) (err error) {
	h.ctx.WithField("registration", reg).Debug("New registration request")
	defer ensureAckNack(an, nil, &err)
	stats.MarkMeter("handler.registration.in")

	hreg, ok := reg.(HRegistration)
	if !ok {
		stats.MarkMeter("handler.registration.invalid")
		return errors.New(errors.Structural, "Not a Handler registration")
	}

	if err = h.devices.StorePersonalized(hreg); err != nil {
		return errors.New(errors.Operational, err)
	}

	return sub.Subscribe(brokerRegistration{
		recipient: h.broker,
		appEUI:    hreg.AppEUI(),
		devEUI:    hreg.DevEUI(),
		nwkSKey:   hreg.NwkSKey(),
	})
}

// HandleUp implements the core.Component interface
func (h component) HandleUp(data []byte, an AckNacker, up Adapter) (err error) {
	// Make sure we don't forget the AckNacker
	var ack Packet
	defer ensureAckNack(an, &ack, &err)
	stats.MarkMeter("handler.uplink.in")

	itf, err := UnmarshalPacket(data)
	if err != nil {
		stats.MarkMeter("handler.uplink.invalid")
		return errors.New(errors.Structural, err)
	}

	switch itf.(type) {
	case HPacket:
		stats.MarkMeter("handler.uplink.data")

		// 0. Retrieve the handler packet
		packet := itf.(HPacket)
		appEUI := packet.AppEUI()
		devEUI := packet.DevEUI()

		// 1. Lookup for the associated AppSKey + Recipient
		h.ctx.WithField("appEUI", appEUI).WithField("devEUI", devEUI).Debug("Perform lookup")
		entry, err := h.devices.Lookup(appEUI, devEUI)
		if err != nil {
			return err
		}

		// 2. Prepare a channel to receive the response from the consumer
		chresp := make(chan interface{})

		// 3. Create a "bundle" which holds info waiting for other related packets
		var bundleID [20]byte // AppEUI(8) | DevEUI(8) | FCnt
		buf := new(bytes.Buffer)
		binary.Write(buf, binary.BigEndian, appEUI[:])
		binary.Write(buf, binary.BigEndian, devEUI[:])
		binary.Write(buf, binary.BigEndian, packet.FCnt())
		data := buf.Bytes()
		if len(data) != 20 {
			return errors.New(errors.Structural, "Unable to generate bundleID")
		}
		copy(bundleID[:], data[:])

		// 4. Send the actual bundle to the consumer
		ctx := h.ctx.WithField("BundleID", bundleID)
		ctx.Debug("Define new bundle")
		h.set <- bundle{
			ID:      bundleID,
			Packet:  packet,
			Entry:   entry,
			Adapter: up,
			Chresp:  chresp,
			Time:    time.Now(),
		}

		// 5. Wait for the response. Could be an error, a packet or nothing.
		// We'll respond to a maximum of one node. The handler will use the
		// rssi + gateway's duty cycle to select to best fit.
		// All other channels will get a nil response.
		// If there's an error, all channels get the error.
		resp := <-chresp
		switch resp.(type) {
		case BPacket:
			stats.MarkMeter("handler.uplink.ack.with_response")
			stats.MarkMeter("handler.downlink.out")
			ctx.Debug("Received response with packet. Sending Ack")
			ack = resp.(Packet)
		case error:
			stats.MarkMeter("handler.uplink.error")
			ctx.WithError(resp.(error)).Warn("Received errored response. Sending Ack")
			return errors.New(errors.Operational, resp.(error))
		default:
			stats.MarkMeter("handler.uplink.ack.without_response")
			ctx.Debug("Received empty response. Sending empty Ack")
		}

		return nil
	case JPacket:
		stats.MarkMeter("handler.uplink.join_request")
		return errors.New(errors.Implementation, "Join Request not yet implemented")
	default:
		stats.MarkMeter("handler.uplink.unknown")
		return errors.New(errors.Implementation, "Unhandled packet type")
	}
}

// consumeBundles processes list of bundle generated overtime, decrypt the underlying packet,
// deduplicate them, and send a single enhanced packet to the upadapter for further processing.
func (h component) consumeBundles(chbundle <-chan []bundle) {
	ctx := h.ctx.WithField("goroutine", "bundle consumer")
	ctx.Debug("Starting bundle consumer")

browseBundles:
	for bundles := range chbundle {
		ctx.WithField("BundleID", bundles[0].ID).Debug("Consume new bundle")
		var metadata []Metadata
		var payload []byte
		var firstTime time.Time

		if len(bundles) < 1 {
			continue browseBundles
		}
		b := bundles[0]

		computer, err := newScoreComputer(b.Packet.Metadata().Datr)
		if err != nil {
			go h.abortConsume(err, bundles)
			continue browseBundles
		}

		stats.UpdateHistogram("handler.uplink.duplicate.count", int64(len(bundles)))

		for i, bundle := range bundles {
			// We only decrypt the payload of the first bundle's packet.
			// We assume all the other to be equal and we'll merely collect
			// metadata from other bundle.
			if i == 0 {
				var err error
				payload, err = bundle.Packet.Payload(bundle.Entry.AppSKey)
				if err != nil {
					go h.abortConsume(err, bundles)
					continue browseBundles
				}
				firstTime = bundle.Time
				stats.MarkMeter("handler.uplink.in.unique")
			} else {
				diff := bundle.Time.Sub(firstTime).Nanoseconds()
				stats.UpdateHistogram("handler.uplink.duplicate.delay", diff/1000)
			}

			// Append metadata for each of them
			metadata = append(metadata, bundle.Packet.Metadata())
			computer.Update(i, bundle.Packet.Metadata())
		}

		// Then create an application-level packet
		packet, err := NewAPacket(b.Packet.AppEUI(), b.Packet.DevEUI(), payload, metadata)
		if err != nil {
			go h.abortConsume(err, bundles)
			continue browseBundles
		}

		// And send it to the wild open
		// we don't expect a response from the adapter, end of the chain.
		recipient, err := b.Adapter.GetRecipient(b.Entry.Recipient)
		if err != nil {
			go h.abortConsume(err, bundles)
			continue browseBundles
		}

		_, err = b.Adapter.Send(packet, recipient)
		if err != nil {
			go h.abortConsume(err, bundles)
			continue browseBundles
		}
		stats.MarkMeter("handler.uplink.out")

		// Now handle the downlink and respond to node
		best := computer.Get()
		down, err := h.packets.Pull(b.Packet.AppEUI(), b.Packet.DevEUI())
		if err != nil && err.(errors.Failure).Nature != errors.NotFound {
			go h.abortConsume(err, bundles)
			continue browseBundles
		}
		for i, bundle := range bundles {
			if best != nil && best.ID == i && down != nil && err == nil {
				stats.MarkMeter("handler.downlink.pull")
				bpacket, err := h.buildDownlink(down, bundle.Packet, bundle.Entry, best.IsRX2)
				if err != nil {
					go h.abortConsume(errors.New(errors.Structural, err), bundles)
					continue browseBundles
				}
				bundle.Chresp <- bpacket
			} else {
				bundle.Chresp <- nil
			}
		}
	}
}

// Abort consume forward the given error to all bundle recipients
func (h component) abortConsume(fault error, bundles []bundle) {
	err := errors.New(errors.Structural, fault)
	stats.MarkMeter("handler.uplink.invalid")
	h.ctx.WithError(err).Debug("Unable to consume bundle")
	for _, bundle := range bundles {
		bundle.Chresp <- err
	}
}

// constructs a downlink packet from something we pulled from the gathered downlink, and, the actual
// uplink.
func (h component) buildDownlink(down APacket, up HPacket, entry devEntry, isRX2 bool) (BPacket, error) {
	fcnt := up.FCnt() + 1
	macPayload := lorawan.NewMACPayload(false)
	macPayload.FHDR = lorawan.FHDR{
		DevAddr: entry.DevAddr,
		FCnt:    fcnt,
	}
	macPayload.FPort = 1
	macPayload.FRMPayload = []lorawan.Payload{&lorawan.DataPayload{
		Bytes: down.Payload(),
	}}

	if err := macPayload.EncryptFRMPayload(entry.AppSKey); err != nil {
		return nil, err
	}

	payload := lorawan.NewPHYPayload(false)
	payload.MHDR = lorawan.MHDR{
		MType: lorawan.UnconfirmedDataDown, // TODO Handle Confirmed data down
		Major: lorawan.LoRaWANR1,
	}
	payload.MACPayload = macPayload

	data, err := payload.MarshalBinary()
	if err != nil {
		return nil, err
	}
	pmetadata := up.Metadata()

	if pmetadata.Tmst == nil || pmetadata.Freq == nil || pmetadata.Codr == nil || pmetadata.Datr == nil {
		return nil, errors.New(errors.Structural, "Missing mandatory metadata in uplink packet")
	}

	metadata := Metadata{
		Freq: pmetadata.Freq,
		Codr: pmetadata.Codr,
		Datr: pmetadata.Datr,
		Size: pointer.Uint(uint(len(data))),
		Tmst: pointer.Uint(*pmetadata.Tmst + 1000),
	}

	if isRX2 { // Should we reply on RX2, metadata aren't the same
		// TODO Handle different regions with non hard-coded values
		metadata.Freq = pointer.Float64(869.5)
		metadata.Datr = pointer.String("SF9BW125")
		metadata.Tmst = pointer.Uint(*pmetadata.Tmst + 2000)
	}

	return NewBPacket(payload, metadata)
}

// consumeSet gathers new incoming bundles which possess the same id (i.e. appEUI & devEUI & Fcnt)
// It then flushes them once a given delay has passed since the reception of the first bundle.
func (h component) consumeSet(chbundles chan<- []bundle, chset <-chan bundle) {
	ctx := h.ctx.WithField("goroutine", "set consumer")
	ctx.Debug("Starting packets buffering")

	// NOTE Processed is likely to grow quickly. One has to define a more efficient data stucture
	// with a ttl for each entry. Processed is merely there to avoid late packets from being
	// processed again. The TTL could be only of several seconds or minutes.
	processed := make(map[[16]byte][]byte) // AppEUI | DevEUI | FCnt -> hasBeenProcessed ?
	buffers := make(map[[20]byte][]bundle) // AppEUI | DevEUI | FCnt ->  buffered bundles
	alarm := make(chan [20]byte)           // Communication channel with subsequent alarms

	for {
		select {
		case id := <-alarm:
			// Get all bundles
			bundles := buffers[id]
			delete(buffers, id)

			// Register the last processed entry
			var pid [16]byte
			copy(pid[:], id[:16])
			processed[pid] = id[16:]

			// Actually send the bundle to the be processed
			go func(bundles []bundle) { chbundles <- bundles }(bundles)
			ctx.WithField("BundleID", id).Debug("Consuming collected bundles")
		case b := <-chset:
			ctx = ctx.WithField("BundleID", b.ID)

			// Check if bundle has already been processed
			var pid [16]byte
			copy(pid[:], b.ID[:16])
			if reflect.DeepEqual(processed[pid], b.ID[16:]) {
				ctx.Debug("Reject already processed bundle")
				go func(b bundle) {
					b.Chresp <- errors.New(errors.Behavioural, "Already processed")
				}(b)
				continue
			}

			// Add the bundle to the stack, and set the alarm if its the first
			bundles := append(buffers[b.ID], b)
			if len(bundles) == 1 {
				go setAlarm(alarm, b.ID, bufferDelay)
				ctx.Debug("Buffering started -> new alarm set")
			}
			buffers[b.ID] = bundles
		}
	}
}

// setAlarm will trigger a message on the given channel after the given delay
func setAlarm(alarm chan<- [20]byte, id [20]byte, delay time.Duration) {
	<-time.After(delay)
	alarm <- id
}

// HandleDown implements the core.Component interface
func (h component) HandleDown(data []byte, an AckNacker, down Adapter) (err error) {
	// Make sure we don't forget the AckNacker
	var ack Packet
	defer ensureAckNack(an, &ack, &err)
	stats.MarkMeter("handler.downlink.in")

	// Unmarshal the given packet and see what gift we get
	itf, err := UnmarshalPacket(data)
	if err != nil {
		stats.MarkMeter("handler.downlink.invalid")
		return errors.New(errors.Structural, err)
	}

	switch itf.(type) {
	case APacket:
		return h.packets.Push(itf.(APacket))
	default:
		stats.MarkMeter("handler.downlink.invalid")
		return errors.New(errors.Implementation, "Unhandled packet type")
	}
}

func ensureAckNack(an AckNacker, ack *Packet, err *error) {
	if err != nil && *err != nil {
		an.Nack(*err)
	} else {
		var p Packet
		if ack != nil {
			p = *ack
		}
		an.Ack(p)
	}
}
<commit_msg>[feature/downlink] Adapt handler to new ScoreComputer<commit_after>// Copyright © 2016 The Things Network
// Use of this source code is governed by the MIT license that can be found in the LICENSE file.

package handler

import (
	"bytes"
	"encoding/binary"
	"reflect"
	"time"

	. "github.com/TheThingsNetwork/ttn/core"
	"github.com/TheThingsNetwork/ttn/core/dutycycle"
	"github.com/TheThingsNetwork/ttn/utils/errors"
	"github.com/TheThingsNetwork/ttn/utils/pointer"
	"github.com/TheThingsNetwork/ttn/utils/stats"
	"github.com/apex/log"
	"github.com/brocaar/lorawan"
)

const bufferDelay time.Duration = time.Millisecond * 300

// component implements the core.Component interface
type component struct {
	broker  JSONRecipient
	ctx     log.Interface
	devices DevStorage
	packets PktStorage
	set     chan<- bundle
}

type bundle struct {
	Adapter Adapter
	Chresp  chan interface{}
	Entry   devEntry
	ID      [20]byte
	Packet  HPacket
	Time    time.Time
}

// New construct a new Handler
func New(devDb DevStorage, pktDb PktStorage, broker JSONRecipient, ctx log.Interface) Handler {
	h := component{
		ctx:     ctx,
		devices: devDb,
		packets: pktDb,
		broker:  broker,
	}

	set := make(chan bundle)
	bundles := make(chan []bundle)

	h.set = set
	go h.consumeBundles(bundles)
	go h.consumeSet(bundles, set)

	return h
}

// Register implements the core.Component interface
func (h component) Register(reg Registration, an AckNacker, sub Subscriber) (err error) {
	h.ctx.WithField("registration", reg).Debug("New registration request")
	defer ensureAckNack(an, nil, &err)
	stats.MarkMeter("handler.registration.in")

	hreg, ok := reg.(HRegistration)
	if !ok {
		stats.MarkMeter("handler.registration.invalid")
		return errors.New(errors.Structural, "Not a Handler registration")
	}

	if err = h.devices.StorePersonalized(hreg); err != nil {
		return errors.New(errors.Operational, err)
	}

	return sub.Subscribe(brokerRegistration{
		recipient: h.broker,
		appEUI:    hreg.AppEUI(),
		devEUI:    hreg.DevEUI(),
		nwkSKey:   hreg.NwkSKey(),
	})
}

// HandleUp implements the core.Component interface
func (h component) HandleUp(data []byte, an AckNacker, up Adapter) (err error) {
	// Make sure we don't forget the AckNacker
	var ack Packet
	defer ensureAckNack(an, &ack, &err)
	stats.MarkMeter("handler.uplink.in")

	itf, err := UnmarshalPacket(data)
	if err != nil {
		stats.MarkMeter("handler.uplink.invalid")
		return errors.New(errors.Structural, err)
	}

	switch itf.(type) {
	case HPacket:
		stats.MarkMeter("handler.uplink.data")

		// 0. Retrieve the handler packet
		packet := itf.(HPacket)
		appEUI := packet.AppEUI()
		devEUI := packet.DevEUI()

		// 1. Lookup for the associated AppSKey + Recipient
		h.ctx.WithField("appEUI", appEUI).WithField("devEUI", devEUI).Debug("Perform lookup")
		entry, err := h.devices.Lookup(appEUI, devEUI)
		if err != nil {
			return err
		}

		// 2. Prepare a channel to receive the response from the consumer
		chresp := make(chan interface{})

		// 3. Create a "bundle" which holds info waiting for other related packets
		var bundleID [20]byte // AppEUI(8) | DevEUI(8) | FCnt
		buf := new(bytes.Buffer)
		binary.Write(buf, binary.BigEndian, appEUI[:])
		binary.Write(buf, binary.BigEndian, devEUI[:])
		binary.Write(buf, binary.BigEndian, packet.FCnt())
		data := buf.Bytes()
		if len(data) != 20 {
			return errors.New(errors.Structural, "Unable to generate bundleID")
		}
		copy(bundleID[:], data[:])

		// 4. Send the actual bundle to the consumer
		ctx := h.ctx.WithField("BundleID", bundleID)
		ctx.Debug("Define new bundle")
		h.set <- bundle{
			ID:      bundleID,
			Packet:  packet,
			Entry:   entry,
			Adapter: up,
			Chresp:  chresp,
			Time:    time.Now(),
		}

		// 5. Wait for the response. Could be an error, a packet or nothing.
		// We'll respond to a maximum of one node. The handler will use the
		// rssi + gateway's duty cycle to select to best fit.
		// All other channels will get a nil response.
		// If there's an error, all channels get the error.
		resp := <-chresp
		switch resp.(type) {
		case BPacket:
			stats.MarkMeter("handler.uplink.ack.with_response")
			stats.MarkMeter("handler.downlink.out")
			ctx.Debug("Received response with packet. Sending Ack")
			ack = resp.(Packet)
		case error:
			stats.MarkMeter("handler.uplink.error")
			ctx.WithError(resp.(error)).Warn("Received errored response. Sending Ack")
			return errors.New(errors.Operational, resp.(error))
		default:
			stats.MarkMeter("handler.uplink.ack.without_response")
			ctx.Debug("Received empty response. Sending empty Ack")
		}

		return nil
	case JPacket:
		stats.MarkMeter("handler.uplink.join_request")
		return errors.New(errors.Implementation, "Join Request not yet implemented")
	default:
		stats.MarkMeter("handler.uplink.unknown")
		return errors.New(errors.Implementation, "Unhandled packet type")
	}
}

// consumeBundles processes list of bundle generated overtime, decrypt the underlying packet,
// deduplicate them, and send a single enhanced packet to the upadapter for further processing.
func (h component) consumeBundles(chbundle <-chan []bundle) {
	ctx := h.ctx.WithField("goroutine", "bundle consumer")
	ctx.Debug("Starting bundle consumer")

browseBundles:
	for bundles := range chbundle {
		ctx.WithField("BundleID", bundles[0].ID).Debug("Consume new bundle")
		var metadata []Metadata
		var payload []byte
		var firstTime time.Time

		if len(bundles) < 1 {
			continue browseBundles
		}
		b := bundles[0]

		computer, scores, err := dutycycle.NewScoreComputer(b.Packet.Metadata().Datr)
		if err != nil {
			go h.abortConsume(err, bundles)
			continue browseBundles
		}

		stats.UpdateHistogram("handler.uplink.duplicate.count", int64(len(bundles)))

		for i, bundle := range bundles {
			// We only decrypt the payload of the first bundle's packet.
			// We assume all the other to be equal and we'll merely collect
			// metadata from other bundle.
			if i == 0 {
				var err error
				payload, err = bundle.Packet.Payload(bundle.Entry.AppSKey)
				if err != nil {
					go h.abortConsume(err, bundles)
					continue browseBundles
				}
				firstTime = bundle.Time
				stats.MarkMeter("handler.uplink.in.unique")
			} else {
				diff := bundle.Time.Sub(firstTime).Nanoseconds()
				stats.UpdateHistogram("handler.uplink.duplicate.delay", diff/1000)
			}

			// Append metadata for each of them
			metadata = append(metadata, bundle.Packet.Metadata())
			scores = computer.Update(scores, i, bundle.Packet.Metadata())
		}

		// Then create an application-level packet
		packet, err := NewAPacket(b.Packet.AppEUI(), b.Packet.DevEUI(), payload, metadata)
		if err != nil {
			go h.abortConsume(err, bundles)
			continue browseBundles
		}

		// And send it to the wild open
		// we don't expect a response from the adapter, end of the chain.
		recipient, err := b.Adapter.GetRecipient(b.Entry.Recipient)
		if err != nil {
			go h.abortConsume(err, bundles)
			continue browseBundles
		}

		_, err = b.Adapter.Send(packet, recipient)
		if err != nil {
			go h.abortConsume(err, bundles)
			continue browseBundles
		}
		stats.MarkMeter("handler.uplink.out")

		// Now handle the downlink and respond to node
		best := computer.Get(scores)
		down, err := h.packets.Pull(b.Packet.AppEUI(), b.Packet.DevEUI())
		if err != nil && err.(errors.Failure).Nature != errors.NotFound {
			go h.abortConsume(err, bundles)
			continue browseBundles
		}
		for i, bundle := range bundles {
			if best != nil && best.ID == i && down != nil && err == nil {
				stats.MarkMeter("handler.downlink.pull")
				bpacket, err := h.buildDownlink(down, bundle.Packet, bundle.Entry, best.IsRX2)
				if err != nil {
					go h.abortConsume(errors.New(errors.Structural, err), bundles)
					continue browseBundles
				}
				bundle.Chresp <- bpacket
			} else {
				bundle.Chresp <- nil
			}
		}
	}
}

// Abort consume forward the given error to all bundle recipients
func (h component) abortConsume(fault error, bundles []bundle) {
	err := errors.New(errors.Structural, fault)
	stats.MarkMeter("handler.uplink.invalid")
	h.ctx.WithError(err).Debug("Unable to consume bundle")
	for _, bundle := range bundles {
		bundle.Chresp <- err
	}
}

// constructs a downlink packet from something we pulled from the gathered downlink, and, the actual
// uplink.
func (h component) buildDownlink(down APacket, up HPacket, entry devEntry, isRX2 bool) (BPacket, error) {
	fcnt := up.FCnt() + 1
	macPayload := lorawan.NewMACPayload(false)
	macPayload.FHDR = lorawan.FHDR{
		DevAddr: entry.DevAddr,
		FCnt:    fcnt,
	}
	macPayload.FPort = 1
	macPayload.FRMPayload = []lorawan.Payload{&lorawan.DataPayload{
		Bytes: down.Payload(),
	}}

	if err := macPayload.EncryptFRMPayload(entry.AppSKey); err != nil {
		return nil, err
	}

	payload := lorawan.NewPHYPayload(false)
	payload.MHDR = lorawan.MHDR{
		MType: lorawan.UnconfirmedDataDown, // TODO Handle Confirmed data down
		Major: lorawan.LoRaWANR1,
	}
	payload.MACPayload = macPayload

	data, err := payload.MarshalBinary()
	if err != nil {
		return nil, err
	}
	pmetadata := up.Metadata()

	if pmetadata.Tmst == nil || pmetadata.Freq == nil || pmetadata.Codr == nil || pmetadata.Datr == nil {
		return nil, errors.New(errors.Structural, "Missing mandatory metadata in uplink packet")
	}

	metadata := Metadata{
		Freq: pmetadata.Freq,
		Codr: pmetadata.Codr,
		Datr: pmetadata.Datr,
		Size: pointer.Uint(uint(len(data))),
		Tmst: pointer.Uint(*pmetadata.Tmst + 1000),
	}

	if isRX2 { // Should we reply on RX2, metadata aren't the same
		// TODO Handle different regions with non hard-coded values
		metadata.Freq = pointer.Float64(869.5)
		metadata.Datr = pointer.String("SF9BW125")
		metadata.Tmst = pointer.Uint(*pmetadata.Tmst + 2000)
	}

	return NewBPacket(payload, metadata)
}

// consumeSet gathers new incoming bundles which possess the same id (i.e. appEUI & devEUI & Fcnt)
// It then flushes them once a given delay has passed since the reception of the first bundle.
func (h component) consumeSet(chbundles chan<- []bundle, chset <-chan bundle) {
	ctx := h.ctx.WithField("goroutine", "set consumer")
	ctx.Debug("Starting packets buffering")

	// NOTE Processed is likely to grow quickly. One has to define a more efficient data stucture
	// with a ttl for each entry. Processed is merely there to avoid late packets from being
	// processed again. The TTL could be only of several seconds or minutes.
	processed := make(map[[16]byte][]byte) // AppEUI | DevEUI | FCnt -> hasBeenProcessed ?
	buffers := make(map[[20]byte][]bundle) // AppEUI | DevEUI | FCnt ->  buffered bundles
	alarm := make(chan [20]byte)           // Communication channel with subsequent alarms

	for {
		select {
		case id := <-alarm:
			// Get all bundles
			bundles := buffers[id]
			delete(buffers, id)

			// Register the last processed entry
			var pid [16]byte
			copy(pid[:], id[:16])
			processed[pid] = id[16:]

			// Actually send the bundle to the be processed
			go func(bundles []bundle) { chbundles <- bundles }(bundles)
			ctx.WithField("BundleID", id).Debug("Consuming collected bundles")
		case b := <-chset:
			ctx = ctx.WithField("BundleID", b.ID)

			// Check if bundle has already been processed
			var pid [16]byte
			copy(pid[:], b.ID[:16])
			if reflect.DeepEqual(processed[pid], b.ID[16:]) {
				ctx.Debug("Reject already processed bundle")
				go func(b bundle) {
					b.Chresp <- errors.New(errors.Behavioural, "Already processed")
				}(b)
				continue
			}

			// Add the bundle to the stack, and set the alarm if its the first
			bundles := append(buffers[b.ID], b)
			if len(bundles) == 1 {
				go setAlarm(alarm, b.ID, bufferDelay)
				ctx.Debug("Buffering started -> new alarm set")
			}
			buffers[b.ID] = bundles
		}
	}
}

// setAlarm will trigger a message on the given channel after the given delay
func setAlarm(alarm chan<- [20]byte, id [20]byte, delay time.Duration) {
	<-time.After(delay)
	alarm <- id
}

// HandleDown implements the core.Component interface
func (h component) HandleDown(data []byte, an AckNacker, down Adapter) (err error) {
	// Make sure we don't forget the AckNacker
	var ack Packet
	defer ensureAckNack(an, &ack, &err)
	stats.MarkMeter("handler.downlink.in")

	// Unmarshal the given packet and see what gift we get
	itf, err := UnmarshalPacket(data)
	if err != nil {
		stats.MarkMeter("handler.downlink.invalid")
		return errors.New(errors.Structural, err)
	}

	switch itf.(type) {
	case APacket:
		return h.packets.Push(itf.(APacket))
	default:
		stats.MarkMeter("handler.downlink.invalid")
		return errors.New(errors.Implementation, "Unhandled packet type")
	}
}

func ensureAckNack(an AckNacker, ack *Packet, err *error) {
	if err != nil && *err != nil {
		an.Nack(*err)
	} else {
		var p Packet
		if ack != nil {
			p = *ack
		}
		an.Ack(p)
	}
}
<|endoftext|>


<commit_before>// Copyright (c) 2015 Marin Atanasov Nikolov <dnaeon@gmail.com>
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions
// are met:
// 1. Redistributions of source code must retain the above copyright
//    notice, this list of conditions and the following disclaimer
//    in this position and unchanged.
// 2. Redistributions in binary form must reproduce the above copyright
//    notice, this list of conditions and the following disclaimer in the
//    documentation and/or other materials provided with the distribution.
//
// THIS SOFTWARE IS PROVIDED BY THE AUTHOR(S) ``AS IS'' AND ANY EXPRESS OR
// IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
// OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
// IN NO EVENT SHALL THE AUTHOR(S) BE LIABLE FOR ANY DIRECT, INDIRECT,
// INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT
// NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
// DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
// THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF
// THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

package recorder

import (
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"net/http/httptest"
	"net/url"
	"os"

	"github.com/dnaeon/go-vcr/cassette"
)

// Recorder states
const (
	ModeRecording = iota
	ModeReplaying
)

type Recorder struct {
	// Operating mode of the recorder
	mode int

	// HTTP server used to mock requests
	server *httptest.Server

	// Cassette used by the recorder
	cassette *cassette.Cassette

	// Transport that can be used by clients to inject
	Transport *http.Transport
}

// Proxies client requests to their original destination
func requestHandler(r *http.Request, c *cassette.Cassette, mode int) (*cassette.Interaction, error) {
	// Return interaction from cassette if in replay mode
	if mode == ModeReplaying {
		return c.GetInteraction(r)
	}

	// Else, perform client request to their original
	// destination and record interactions
	body := ioutil.NopCloser(r.Body)
	req, err := http.NewRequest(r.Method, r.URL.String(), body)
	if err != nil {
		return nil, err
	}

	req.Header = r.Header
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return nil, err
	}

	// Record the interaction and add it to the cassette
	reqBody, err := ioutil.ReadAll(req.Body)
	if err != nil {
		return nil, err
	}

	respBody, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	// Add interaction to cassette
	interaction := &cassette.Interaction{
		Request: cassette.Request{
			Body:    string(reqBody),
			Headers: req.Header,
			URL:     req.URL.String(),
			Method:  req.Method,
		},
		Response: cassette.Response{
			Body:    string(respBody),
			Headers: resp.Header,
			Status:  resp.Status,
			Code:    resp.StatusCode,
		},
	}
	c.AddInteraction(interaction)

	return interaction, nil
}

// Creates a new recorder
func New(cassetteName string) (*Recorder, error) {
	var mode int
	var c *cassette.Cassette
	cassetteFile := fmt.Sprintf("%s.yaml", cassetteName)

	// Depending on whether the cassette file exists or not we
	// either create a new empty cassette or load from file
	if _, err := os.Stat(cassetteFile); os.IsNotExist(err) {
		// Create new cassette and enter in recording mode
		c = cassette.New(cassetteName)
		mode = ModeRecording
	} else {
		// Load cassette from file and enter replay mode
		c, err = cassette.Load(cassetteName)
		if err != nil {
			return nil, err
		}
		mode = ModeReplaying
	}

	// Handler for client requests
	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// Pass cassette and mode to handler, so that interactions can be
		// retrieved or recorded depending on the current recorder mode
		interaction, err := requestHandler(r, c, mode)

		if err != nil {
			log.Fatalf("Failed to process request for URL %s: %s", r.URL, err)
		}

		w.WriteHeader(interaction.Response.Code)
		fmt.Fprintln(w, interaction.Response.Body)
	})

	// HTTP server used to mock requests
	server := httptest.NewServer(handler)

	// A proxy function which routes all requests through our HTTP server
	// Can be used by clients to inject into their own transports
	proxyUrl, err := url.Parse(server.URL)
	if err != nil {
		return nil, err
	}

	// A transport which can be used by clients to inject
	transport := &http.Transport{
		Proxy: http.ProxyURL(proxyUrl),
	}

	r := &Recorder{
		mode:      mode,
		server:    server,
		cassette:  c,
		Transport: transport,
	}

	return r, nil
}

// Stops the recorder
func (r *Recorder) Stop() error {
	r.server.Close()

	if r.mode == ModeRecording {
		if err := r.cassette.Save(); err != nil {
			return err
		}
	}

	return nil
}
<commit_msg>Panic instead of os.Exit(1) on missing interaction<commit_after>// Copyright (c) 2015 Marin Atanasov Nikolov <dnaeon@gmail.com>
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions
// are met:
// 1. Redistributions of source code must retain the above copyright
//    notice, this list of conditions and the following disclaimer
//    in this position and unchanged.
// 2. Redistributions in binary form must reproduce the above copyright
//    notice, this list of conditions and the following disclaimer in the
//    documentation and/or other materials provided with the distribution.
//
// THIS SOFTWARE IS PROVIDED BY THE AUTHOR(S) ``AS IS'' AND ANY EXPRESS OR
// IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
// OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
// IN NO EVENT SHALL THE AUTHOR(S) BE LIABLE FOR ANY DIRECT, INDIRECT,
// INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT
// NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
// DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
// THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF
// THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

package recorder

import (
	"fmt"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"net/url"
	"os"

	"github.com/dnaeon/go-vcr/cassette"
)

// Recorder states
const (
	ModeRecording = iota
	ModeReplaying
)

type Recorder struct {
	// Operating mode of the recorder
	mode int

	// HTTP server used to mock requests
	server *httptest.Server

	// Cassette used by the recorder
	cassette *cassette.Cassette

	// Transport that can be used by clients to inject
	Transport *http.Transport
}

// Proxies client requests to their original destination
func requestHandler(r *http.Request, c *cassette.Cassette, mode int) (*cassette.Interaction, error) {
	// Return interaction from cassette if in replay mode
	if mode == ModeReplaying {
		return c.GetInteraction(r)
	}

	// Else, perform client request to their original
	// destination and record interactions
	body := ioutil.NopCloser(r.Body)
	req, err := http.NewRequest(r.Method, r.URL.String(), body)
	if err != nil {
		return nil, err
	}

	req.Header = r.Header
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return nil, err
	}

	// Record the interaction and add it to the cassette
	reqBody, err := ioutil.ReadAll(req.Body)
	if err != nil {
		return nil, err
	}

	respBody, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	// Add interaction to cassette
	interaction := &cassette.Interaction{
		Request: cassette.Request{
			Body:    string(reqBody),
			Headers: req.Header,
			URL:     req.URL.String(),
			Method:  req.Method,
		},
		Response: cassette.Response{
			Body:    string(respBody),
			Headers: resp.Header,
			Status:  resp.Status,
			Code:    resp.StatusCode,
		},
	}
	c.AddInteraction(interaction)

	return interaction, nil
}

// Creates a new recorder
func New(cassetteName string) (*Recorder, error) {
	var mode int
	var c *cassette.Cassette
	cassetteFile := fmt.Sprintf("%s.yaml", cassetteName)

	// Depending on whether the cassette file exists or not we
	// either create a new empty cassette or load from file
	if _, err := os.Stat(cassetteFile); os.IsNotExist(err) {
		// Create new cassette and enter in recording mode
		c = cassette.New(cassetteName)
		mode = ModeRecording
	} else {
		// Load cassette from file and enter replay mode
		c, err = cassette.Load(cassetteName)
		if err != nil {
			return nil, err
		}
		mode = ModeReplaying
	}

	// Handler for client requests
	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// Pass cassette and mode to handler, so that interactions can be
		// retrieved or recorded depending on the current recorder mode
		interaction, err := requestHandler(r, c, mode)

		if err != nil {
			panic(fmt.Errorf("Failed to process request for URL %s: %s", r.URL, err))
		}

		w.WriteHeader(interaction.Response.Code)
		fmt.Fprintln(w, interaction.Response.Body)
	})

	// HTTP server used to mock requests
	server := httptest.NewServer(handler)

	// A proxy function which routes all requests through our HTTP server
	// Can be used by clients to inject into their own transports
	proxyUrl, err := url.Parse(server.URL)
	if err != nil {
		return nil, err
	}

	// A transport which can be used by clients to inject
	transport := &http.Transport{
		Proxy: http.ProxyURL(proxyUrl),
	}

	r := &Recorder{
		mode:      mode,
		server:    server,
		cassette:  c,
		Transport: transport,
	}

	return r, nil
}

// Stops the recorder
func (r *Recorder) Stop() error {
	r.server.Close()

	if r.mode == ModeRecording {
		if err := r.cassette.Save(); err != nil {
			return err
		}
	}

	return nil
}
<|endoftext|>