<commit_before>package main

import (
	"bufio"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path"
	"path/filepath"

	"github.com/jessevdk/go-flags"
	"github.com/sjoerdsimons/fakemachine"

	"gopkg.in/yaml.v2"
)

func CleanPathAt(path, at string) string {
	if filepath.IsAbs(path) {
		return filepath.Clean(path)
	}

	return filepath.Join(at, path)
}

func CleanPath(path string) string {
	cwd, _ := os.Getwd()
	return CleanPathAt(path, cwd)
}

func CopyFile(src, dst string, mode os.FileMode) error {
	in, err := os.Open(src)
	if err != nil {
		return err
	}
	defer in.Close()
	tmp, err := ioutil.TempFile(filepath.Dir(dst), "")
	if err != nil {
		return err
	}
	_, err = io.Copy(tmp, in)
	if err != nil {
		tmp.Close()
		os.Remove(tmp.Name())
		return err
	}
	if err = tmp.Close(); err != nil {
		os.Remove(tmp.Name())
		return err
	}
	if err = os.Chmod(tmp.Name(), mode); err != nil {
		os.Remove(tmp.Name())
		return err
	}
	return os.Rename(tmp.Name(), dst)
}

func CopyTree(sourcetree, desttree string) {
	fmt.Printf("Overlaying %s on %s\n", sourcetree, desttree)
	walker := func(p string, info os.FileInfo, err error) error {

		if err != nil {
			return err
		}

		suffix, _ := filepath.Rel(sourcetree, p)
		target := path.Join(desttree, suffix)
		switch info.Mode() & os.ModeType {
		case 0:
			fmt.Printf("F> %s\n", p)
			CopyFile(p, target, info.Mode())
		case os.ModeDir:
			fmt.Printf("D> %s -> %s\n", p, target)
			os.Mkdir(target, info.Mode())
		case os.ModeSymlink:
			link, err := os.Readlink(p)
			if err != nil {
				log.Panic("Failed to read symlink %s: %v", suffix, err)
			}
			os.Symlink(link, target)
		default:
			log.Panicf("Not handled /%s %v", suffix, info.Mode())
		}

		return nil
	}

	filepath.Walk(sourcetree, walker)
}

type QemuHelper struct {
	qemusrc    string
	qemutarget string
}

func NewQemuHelper(context YaibContext) QemuHelper {
	q := QemuHelper{}

	switch context.Architecture {
	case "armhf", "armel", "arm":
		q.qemusrc = "/usr/bin/qemu-arm-static"
	case "arm64":
		q.qemusrc = "/usr/bin/qemu-aarch64-static"
	case "amd64":
		/* Dummy, no qemu */
	default:
		log.Panicf("Don't know qemu for Architecture %s", context.Architecture)
	}

	if q.qemusrc != "" {
		q.qemutarget = path.Join(context.rootdir, q.qemusrc)
	}

	return q
}

func (q QemuHelper) Setup() error {
	if q.qemusrc == "" {
		return nil
	}
	return CopyFile(q.qemusrc, q.qemutarget, 755)
}

func (q QemuHelper) Cleanup() {
	if q.qemusrc != "" {
		os.Remove(q.qemutarget)
	}
}

func RunCommand(label, command string, arg ...string) error {
	cmd := exec.Command(command, arg...)

	output, _ := cmd.StdoutPipe()
	stderr, _ := cmd.StderrPipe()

	cmd.Start()

	fmt.Printf("Running %s: %s %v\n", label, command, arg)

	scanner := bufio.NewScanner(output)
	for scanner.Scan() {
		fmt.Printf("%s | %s\n", label, scanner.Text())
	}

	reader := bufio.NewReader(stderr)

	for {
		line, _, err := reader.ReadLine()

		if err == io.EOF {
			break
		} else if err != nil {
			fmt.Printf("FAILED: %v\n", err)
			break
		}
		fmt.Printf("%s E | %s\n", label, line)
	}

	err := cmd.Wait()

	return err
}

func RunCommandInChroot(context YaibContext, label, command string, arg ...string) error {
	options := []string{"-q", "-D", context.rootdir, command}
	options = append(options, arg...)

	q := NewQemuHelper(context)
	q.Setup()
	defer q.Cleanup()

	return RunCommand(label, "systemd-nspawn", options...)
}

type YaibContext struct {
	scratchdir   string
	rootdir      string
	artifactdir  string
	image        string
	imageMntDir  string
	recipeDir    string
	Architecture string
}

type Action interface {
	/* FIXME verify should probably be prepare or somesuch */
	Verify(context *YaibContext)
	PreMachine(context *YaibContext, m *fakemachine.Machine, args *[]string)
	PreNoMachine(context *YaibContext)
	Run(context *YaibContext)
	Cleanup(context YaibContext)
	PostMachine(context YaibContext)
}

type BaseAction struct{}

func (b *BaseAction) Verify(context *YaibContext) {}
func (b *BaseAction) PreMachine(context *YaibContext,
	m *fakemachine.Machine,
	args *[]string) {
}
func (b *BaseAction) PreNoMachine(context *YaibContext) {}
func (b *BaseAction) Run(context *YaibContext)          {}
func (b *BaseAction) Cleanup(context YaibContext)       {}
func (b *BaseAction) PostMachine(context YaibContext)   {}

/* the YamlAction just embed the Action interface and implements the
 * UnmarshalYAML function so it can select the concrete implementer of a
 * specific action at unmarshaling time */
type YamlAction struct {
	Action
}

func (y *YamlAction) UnmarshalYAML(unmarshal func(interface{}) error) error {
	var aux struct {
		Action string
	}
	err := unmarshal(&aux)
	if err != nil {
		return err
	}

	switch aux.Action {
	case "debootstrap":
		y.Action = &DebootstrapAction{}
	case "pack":
		y.Action = &PackAction{}
	case "unpack":
		y.Action = &UnpackAction{}
	case "run":
		y.Action = &RunAction{}
	case "apt":
		y.Action = &AptAction{}
	case "overlay":
		y.Action = &OverlayAction{}
	case "setup-image":
		y.Action = &SetupImage{}
	case "raw":
		y.Action = &RawAction{}
	default:
		log.Fatalf("Unknown action: %v", aux.Action)
	}

	unmarshal(y.Action)

	return nil
}

type Recipe struct {
	Architecture string
	Actions      []YamlAction
}

func main() {
	var context YaibContext
	var options struct {
		ArtifactDir   string `long:"artifactdir"`
		InternalImage string `long:"internal-image" hidden:"true"`
	}

	parser := flags.NewParser(&options, flags.Default)
	args, err := parser.Parse()

	if err != nil {
		flagsErr, ok := err.(*flags.Error)
		if ok && flagsErr.Type == flags.ErrHelp {
			os.Exit(0)
		} else {
			fmt.Printf("%v\n", flagsErr)
			os.Exit(1)
		}
	}

	if len(args) != 1 {
		log.Fatal("No recipe given!")
	}

	file := args[0]
	file = CleanPath(file)

	/* If fakemachine is supported the outer fake machine will never use the
	 * scratchdir, so just set it to /scrach as a dummy to prevent the outer
	 * yaib createing a temporary direction */
	if fakemachine.InMachine() || fakemachine.Supported() {
		context.scratchdir = "/scratch"
	} else {
		cwd, _ := os.Getwd()
		context.scratchdir, err = ioutil.TempDir(cwd, ".yaib-")
		defer os.RemoveAll(context.scratchdir)
	}

	context.rootdir = path.Join(context.scratchdir, "root")
	context.image = options.InternalImage
	context.recipeDir = path.Dir(file)

	context.artifactdir = options.ArtifactDir
	if context.artifactdir == "" {
		context.artifactdir, _ = os.Getwd()
	}
	context.artifactdir = CleanPath(context.artifactdir)

	data, err := ioutil.ReadFile(file)
	if err != nil {
		panic(err)
	}

	r := Recipe{}

	err = yaml.Unmarshal(data, &r)
	if err != nil {
		panic(err)
	}

	context.Architecture = r.Architecture

	for _, a := range r.Actions {
		a.Verify(&context)
	}

	if !fakemachine.InMachine() && fakemachine.Supported() {
		m := fakemachine.NewMachine()
		var args []string

		m.AddVolume(context.artifactdir)
		args = append(args, "--artifactdir", context.artifactdir)

		m.AddVolume(context.recipeDir)
		args = append(args, file)

		for _, a := range r.Actions {
			a.PreMachine(&context, m, &args)
		}

		ret := m.RunInMachineWithArgs(args)

		for _, a := range r.Actions {
			a.PostMachine(context)
		}

		os.Exit(ret)
	} else {
		for _, a := range r.Actions {
			a.PreNoMachine(&context)
		}
	}

	for _, a := range r.Actions {
		a.Run(&context)
	}

	for _, a := range r.Actions {
		a.Cleanup(context)
	}

	if !fakemachine.Supported() {
		for _, a := range r.Actions {
			a.PostMachine(context)
		}
	}
}
<commit_msg>Correct logic to determine when running in a fake machine<commit_after>package main

import (
	"bufio"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path"
	"path/filepath"

	"github.com/jessevdk/go-flags"
	"github.com/sjoerdsimons/fakemachine"

	"gopkg.in/yaml.v2"
)

func CleanPathAt(path, at string) string {
	if filepath.IsAbs(path) {
		return filepath.Clean(path)
	}

	return filepath.Join(at, path)
}

func CleanPath(path string) string {
	cwd, _ := os.Getwd()
	return CleanPathAt(path, cwd)
}

func CopyFile(src, dst string, mode os.FileMode) error {
	in, err := os.Open(src)
	if err != nil {
		return err
	}
	defer in.Close()
	tmp, err := ioutil.TempFile(filepath.Dir(dst), "")
	if err != nil {
		return err
	}
	_, err = io.Copy(tmp, in)
	if err != nil {
		tmp.Close()
		os.Remove(tmp.Name())
		return err
	}
	if err = tmp.Close(); err != nil {
		os.Remove(tmp.Name())
		return err
	}
	if err = os.Chmod(tmp.Name(), mode); err != nil {
		os.Remove(tmp.Name())
		return err
	}
	return os.Rename(tmp.Name(), dst)
}

func CopyTree(sourcetree, desttree string) {
	fmt.Printf("Overlaying %s on %s\n", sourcetree, desttree)
	walker := func(p string, info os.FileInfo, err error) error {

		if err != nil {
			return err
		}

		suffix, _ := filepath.Rel(sourcetree, p)
		target := path.Join(desttree, suffix)
		switch info.Mode() & os.ModeType {
		case 0:
			fmt.Printf("F> %s\n", p)
			CopyFile(p, target, info.Mode())
		case os.ModeDir:
			fmt.Printf("D> %s -> %s\n", p, target)
			os.Mkdir(target, info.Mode())
		case os.ModeSymlink:
			link, err := os.Readlink(p)
			if err != nil {
				log.Panic("Failed to read symlink %s: %v", suffix, err)
			}
			os.Symlink(link, target)
		default:
			log.Panicf("Not handled /%s %v", suffix, info.Mode())
		}

		return nil
	}

	filepath.Walk(sourcetree, walker)
}

type QemuHelper struct {
	qemusrc    string
	qemutarget string
}

func NewQemuHelper(context YaibContext) QemuHelper {
	q := QemuHelper{}

	switch context.Architecture {
	case "armhf", "armel", "arm":
		q.qemusrc = "/usr/bin/qemu-arm-static"
	case "arm64":
		q.qemusrc = "/usr/bin/qemu-aarch64-static"
	case "amd64":
		/* Dummy, no qemu */
	default:
		log.Panicf("Don't know qemu for Architecture %s", context.Architecture)
	}

	if q.qemusrc != "" {
		q.qemutarget = path.Join(context.rootdir, q.qemusrc)
	}

	return q
}

func (q QemuHelper) Setup() error {
	if q.qemusrc == "" {
		return nil
	}
	return CopyFile(q.qemusrc, q.qemutarget, 755)
}

func (q QemuHelper) Cleanup() {
	if q.qemusrc != "" {
		os.Remove(q.qemutarget)
	}
}

func RunCommand(label, command string, arg ...string) error {
	cmd := exec.Command(command, arg...)

	output, _ := cmd.StdoutPipe()
	stderr, _ := cmd.StderrPipe()

	cmd.Start()

	fmt.Printf("Running %s: %s %v\n", label, command, arg)

	scanner := bufio.NewScanner(output)
	for scanner.Scan() {
		fmt.Printf("%s | %s\n", label, scanner.Text())
	}

	reader := bufio.NewReader(stderr)

	for {
		line, _, err := reader.ReadLine()

		if err == io.EOF {
			break
		} else if err != nil {
			fmt.Printf("FAILED: %v\n", err)
			break
		}
		fmt.Printf("%s E | %s\n", label, line)
	}

	err := cmd.Wait()

	return err
}

func RunCommandInChroot(context YaibContext, label, command string, arg ...string) error {
	options := []string{"-q", "-D", context.rootdir, command}
	options = append(options, arg...)

	q := NewQemuHelper(context)
	q.Setup()
	defer q.Cleanup()

	return RunCommand(label, "systemd-nspawn", options...)
}

type YaibContext struct {
	scratchdir   string
	rootdir      string
	artifactdir  string
	image        string
	imageMntDir  string
	recipeDir    string
	Architecture string
}

type Action interface {
	/* FIXME verify should probably be prepare or somesuch */
	Verify(context *YaibContext)
	PreMachine(context *YaibContext, m *fakemachine.Machine, args *[]string)
	PreNoMachine(context *YaibContext)
	Run(context *YaibContext)
	Cleanup(context YaibContext)
	PostMachine(context YaibContext)
}

type BaseAction struct{}

func (b *BaseAction) Verify(context *YaibContext) {}
func (b *BaseAction) PreMachine(context *YaibContext,
	m *fakemachine.Machine,
	args *[]string) {
}
func (b *BaseAction) PreNoMachine(context *YaibContext) {}
func (b *BaseAction) Run(context *YaibContext)          {}
func (b *BaseAction) Cleanup(context YaibContext)       {}
func (b *BaseAction) PostMachine(context YaibContext)   {}

/* the YamlAction just embed the Action interface and implements the
 * UnmarshalYAML function so it can select the concrete implementer of a
 * specific action at unmarshaling time */
type YamlAction struct {
	Action
}

func (y *YamlAction) UnmarshalYAML(unmarshal func(interface{}) error) error {
	var aux struct {
		Action string
	}
	err := unmarshal(&aux)
	if err != nil {
		return err
	}

	switch aux.Action {
	case "debootstrap":
		y.Action = &DebootstrapAction{}
	case "pack":
		y.Action = &PackAction{}
	case "unpack":
		y.Action = &UnpackAction{}
	case "run":
		y.Action = &RunAction{}
	case "apt":
		y.Action = &AptAction{}
	case "overlay":
		y.Action = &OverlayAction{}
	case "setup-image":
		y.Action = &SetupImage{}
	case "raw":
		y.Action = &RawAction{}
	default:
		log.Fatalf("Unknown action: %v", aux.Action)
	}

	unmarshal(y.Action)

	return nil
}

type Recipe struct {
	Architecture string
	Actions      []YamlAction
}

func main() {
	var context YaibContext
	var options struct {
		ArtifactDir   string `long:"artifactdir"`
		InternalImage string `long:"internal-image" hidden:"true"`
	}

	parser := flags.NewParser(&options, flags.Default)
	args, err := parser.Parse()

	if err != nil {
		flagsErr, ok := err.(*flags.Error)
		if ok && flagsErr.Type == flags.ErrHelp {
			os.Exit(0)
		} else {
			fmt.Printf("%v\n", flagsErr)
			os.Exit(1)
		}
	}

	if len(args) != 1 {
		log.Fatal("No recipe given!")
	}

	file := args[0]
	file = CleanPath(file)

	/* If fakemachine is supported the outer fake machine will never use the
	 * scratchdir, so just set it to /scrach as a dummy to prevent the outer
	 * yaib createing a temporary direction */
	if fakemachine.InMachine() || fakemachine.Supported() {
		context.scratchdir = "/scratch"
	} else {
		cwd, _ := os.Getwd()
		context.scratchdir, err = ioutil.TempDir(cwd, ".yaib-")
		defer os.RemoveAll(context.scratchdir)
	}

	context.rootdir = path.Join(context.scratchdir, "root")
	context.image = options.InternalImage
	context.recipeDir = path.Dir(file)

	context.artifactdir = options.ArtifactDir
	if context.artifactdir == "" {
		context.artifactdir, _ = os.Getwd()
	}
	context.artifactdir = CleanPath(context.artifactdir)

	data, err := ioutil.ReadFile(file)
	if err != nil {
		panic(err)
	}

	r := Recipe{}

	err = yaml.Unmarshal(data, &r)
	if err != nil {
		panic(err)
	}

	context.Architecture = r.Architecture

	for _, a := range r.Actions {
		a.Verify(&context)
	}

	if !fakemachine.InMachine() && fakemachine.Supported() {
		m := fakemachine.NewMachine()
		var args []string

		m.AddVolume(context.artifactdir)
		args = append(args, "--artifactdir", context.artifactdir)

		m.AddVolume(context.recipeDir)
		args = append(args, file)

		for _, a := range r.Actions {
			a.PreMachine(&context, m, &args)
		}

		ret := m.RunInMachineWithArgs(args)

		for _, a := range r.Actions {
			a.PostMachine(context)
		}

		os.Exit(ret)
	}

	if !fakemachine.InMachine() {
		for _, a := range r.Actions {
			a.PreNoMachine(&context)
		}
	}

	for _, a := range r.Actions {
		a.Run(&context)
	}

	for _, a := range r.Actions {
		a.Cleanup(context)
	}

	if !fakemachine.InMachine() {
		for _, a := range r.Actions {
			a.PostMachine(context)
		}
	}
}
<|endoftext|>


<commit_before>// Copyright (C) 2014 Yasuhiro Matsumoto <mattn.jp@gmail.com>.
//
// Use of this source code is governed by an MIT-style
// license that can be found in the LICENSE file.

package sqlite3

import (
	"database/sql"
	"io/ioutil"
	"os"
	"path"
	"testing"
)

func TestSimpleError(t *testing.T) {
	e := ErrError.Error()
	if e != "SQL logic error or missing database" {
		t.Error("wrong error code:" + e)
	}
}

func TestCorruptDbErrors(t *testing.T) {
	dirName, err := ioutil.TempDir("", "sqlite3")
	if err != nil {
		t.Fatal(err)
	}
	defer os.RemoveAll(dirName)

	dbFileName := path.Join(dirName, "test.db")
	f, err := os.Create(dbFileName)
	if err != nil {
		t.Error(err)
	}
	f.Write([]byte{1, 2, 3, 4, 5})
	f.Close()

	db, err := sql.Open("sqlite3", dbFileName)
	if err == nil {
		_, err = db.Exec("drop table foo")
	}

	sqliteErr := err.(*Error)
	if sqliteErr.Code != ErrNotADB {
		t.Error("wrong error code for corrupted DB")
	}
	if err.Error() == "" {
		t.Error("wrong error string for corrupted DB")
	}
	db.Close()
}

func TestSqlLogicErrors(t *testing.T) {
	dirName, err := ioutil.TempDir("", "sqlite3")
	if err != nil {
		t.Fatal(err)
	}
	defer os.RemoveAll(dirName)

	dbFileName := path.Join(dirName, "test.db")
	db, err := sql.Open("sqlite3", dbFileName)
	if err != nil {
		t.Error(err)
	}
	defer db.Close()

	_, err = db.Exec("CREATE TABLE Foo (id INTEGER PRIMARY KEY)")
	if err != nil {
		t.Error(err)
	}

	const expectedErr = "table Foo already exists"
	_, err = db.Exec("CREATE TABLE Foo (id INTEGER PRIMARY KEY)")
	if err.Error() != expectedErr {
		t.Errorf("Unexpected error: %s, expected %s", err.Error(), expectedErr)
	}

}

func TestExtendedErrorCodes_ForeignKey(t *testing.T) {
	dirName, err := ioutil.TempDir("", "sqlite3-err")
	if err != nil {
		t.Fatal(err)
	}
	defer os.RemoveAll(dirName)

	dbFileName := path.Join(dirName, "test.db")
	db, err := sql.Open("sqlite3", dbFileName)
	if err != nil {
		t.Error(err)
	}
	defer db.Close()

	_, err = db.Exec("PRAGMA foreign_keys=ON;")
	if err != nil {
		t.Errorf("PRAGMA foreign_keys=ON: %v", err)
	}

	_, err = db.Exec(`CREATE TABLE Foo (
		id INTEGER PRIMARY KEY AUTOINCREMENT,
		value INTEGER NOT NULL,
		ref INTEGER NULL REFERENCES Foo (id),
		UNIQUE(value)
	);`)
	if err != nil {
		t.Error(err)
	}

	_, err = db.Exec("INSERT INTO Foo (ref, value) VALUES (100, 100);")
	if err == nil {
		t.Error("No error!")
	} else {
		sqliteErr := err.(Error)
		if sqliteErr.Code != ErrConstraint {
			t.Errorf("Wrong basic error code: %d != %d",
				sqliteErr.Code, ErrConstraint)
		}
		if sqliteErr.ExtendedCode != ErrConstraintForeignKey {
			t.Errorf("Wrong extended error code: %d != %d",
				sqliteErr.ExtendedCode, ErrConstraintForeignKey)
		}
	}

}

func TestExtendedErrorCodes_NotNull(t *testing.T) {
	dirName, err := ioutil.TempDir("", "sqlite3-err")
	if err != nil {
		t.Fatal(err)
	}
	defer os.RemoveAll(dirName)

	dbFileName := path.Join(dirName, "test.db")
	db, err := sql.Open("sqlite3", dbFileName)
	if err != nil {
		t.Error(err)
	}
	defer db.Close()

	_, err = db.Exec("PRAGMA foreign_keys=ON;")
	if err != nil {
		t.Errorf("PRAGMA foreign_keys=ON: %v", err)
	}

	_, err = db.Exec(`CREATE TABLE Foo (
		id INTEGER PRIMARY KEY AUTOINCREMENT,
		value INTEGER NOT NULL,
		ref INTEGER NULL REFERENCES Foo (id),
		UNIQUE(value)
	);`)
	if err != nil {
		t.Error(err)
	}

	res, err := db.Exec("INSERT INTO Foo (value) VALUES (100);")
	if err != nil {
		t.Fatalf("Creating first row: %v", err)
	}

	id, err := res.LastInsertId()
	if err != nil {
		t.Fatalf("Retrieving last insert id: %v", err)
	}

	_, err = db.Exec("INSERT INTO Foo (ref) VALUES (?);", id)
	if err == nil {
		t.Error("No error!")
	} else {
		sqliteErr := err.(Error)
		if sqliteErr.Code != ErrConstraint {
			t.Errorf("Wrong basic error code: %d != %d",
				sqliteErr.Code, ErrConstraint)
		}
		if sqliteErr.ExtendedCode != ErrConstraintNotNull {
			t.Errorf("Wrong extended error code: %d != %d",
				sqliteErr.ExtendedCode, ErrConstraintNotNull)
		}
	}

}

func TestExtendedErrorCodes_Unique(t *testing.T) {
	dirName, err := ioutil.TempDir("", "sqlite3-err")
	if err != nil {
		t.Fatal(err)
	}
	defer os.RemoveAll(dirName)

	dbFileName := path.Join(dirName, "test.db")
	db, err := sql.Open("sqlite3", dbFileName)
	if err != nil {
		t.Error(err)
	}
	defer db.Close()

	_, err = db.Exec("PRAGMA foreign_keys=ON;")
	if err != nil {
		t.Errorf("PRAGMA foreign_keys=ON: %v", err)
	}

	_, err = db.Exec(`CREATE TABLE Foo (
		id INTEGER PRIMARY KEY AUTOINCREMENT,
		value INTEGER NOT NULL,
		ref INTEGER NULL REFERENCES Foo (id),
		UNIQUE(value)
	);`)
	if err != nil {
		t.Error(err)
	}

	res, err := db.Exec("INSERT INTO Foo (value) VALUES (100);")
	if err != nil {
		t.Fatalf("Creating first row: %v", err)
	}

	id, err := res.LastInsertId()
	if err != nil {
		t.Fatalf("Retrieving last insert id: %v", err)
	}

	_, err = db.Exec("INSERT INTO Foo (ref, value) VALUES (?, 100);", id)
	if err == nil {
		t.Error("No error!")
	} else {
		sqliteErr := err.(Error)
		if sqliteErr.Code != ErrConstraint {
			t.Errorf("Wrong basic error code: %d != %d",
				sqliteErr.Code, ErrConstraint)
		}
		if sqliteErr.ExtendedCode != ErrConstraintUnique {
			t.Errorf("Wrong extended error code: %d != %d",
				sqliteErr.ExtendedCode, ErrConstraintUnique)
		}
		extended := sqliteErr.Code.Extend(3).Error()
		expected := "constraint failed"
		if extended != expected {
			t.Errorf("Wrong basic error code: %q != %q",
				extended, expected)
		}
	}

}
<commit_msg>fix test<commit_after>// Copyright (C) 2014 Yasuhiro Matsumoto <mattn.jp@gmail.com>.
//
// Use of this source code is governed by an MIT-style
// license that can be found in the LICENSE file.

package sqlite3

import (
	"database/sql"
	"io/ioutil"
	"os"
	"path"
	"testing"
)

func TestSimpleError(t *testing.T) {
	e := ErrError.Error()
	if e != "SQL logic error or missing database" {
		t.Error("wrong error code:" + e)
	}
}

func TestCorruptDbErrors(t *testing.T) {
	dirName, err := ioutil.TempDir("", "sqlite3")
	if err != nil {
		t.Fatal(err)
	}
	defer os.RemoveAll(dirName)

	dbFileName := path.Join(dirName, "test.db")
	f, err := os.Create(dbFileName)
	if err != nil {
		t.Error(err)
	}
	f.Write([]byte{1, 2, 3, 4, 5})
	f.Close()

	db, err := sql.Open("sqlite3", dbFileName)
	if err == nil {
		_, err = db.Exec("drop table foo")
	}

	sqliteErr := err.(*Error)
	if sqliteErr.Code != ErrNotADB {
		t.Error("wrong error code for corrupted DB")
	}
	if err.Error() == "" {
		t.Error("wrong error string for corrupted DB")
	}
	db.Close()
}

func TestSqlLogicErrors(t *testing.T) {
	dirName, err := ioutil.TempDir("", "sqlite3")
	if err != nil {
		t.Fatal(err)
	}
	defer os.RemoveAll(dirName)

	dbFileName := path.Join(dirName, "test.db")
	db, err := sql.Open("sqlite3", dbFileName)
	if err != nil {
		t.Error(err)
	}
	defer db.Close()

	_, err = db.Exec("CREATE TABLE Foo (id INTEGER PRIMARY KEY)")
	if err != nil {
		t.Error(err)
	}

	const expectedErr = "table Foo already exists"
	_, err = db.Exec("CREATE TABLE Foo (id INTEGER PRIMARY KEY)")
	if err.Error() != expectedErr {
		t.Errorf("Unexpected error: %s, expected %s", err.Error(), expectedErr)
	}

}

func TestExtendedErrorCodes_ForeignKey(t *testing.T) {
	dirName, err := ioutil.TempDir("", "sqlite3-err")
	if err != nil {
		t.Fatal(err)
	}
	defer os.RemoveAll(dirName)

	dbFileName := path.Join(dirName, "test.db")
	db, err := sql.Open("sqlite3", dbFileName)
	if err != nil {
		t.Error(err)
	}
	defer db.Close()

	_, err = db.Exec("PRAGMA foreign_keys=ON;")
	if err != nil {
		t.Errorf("PRAGMA foreign_keys=ON: %v", err)
	}

	_, err = db.Exec(`CREATE TABLE Foo (
		id INTEGER PRIMARY KEY AUTOINCREMENT,
		value INTEGER NOT NULL,
		ref INTEGER NULL REFERENCES Foo (id),
		UNIQUE(value)
	);`)
	if err != nil {
		t.Error(err)
	}

	_, err = db.Exec("INSERT INTO Foo (ref, value) VALUES (100, 100);")
	if err == nil {
		t.Error("No error!")
	} else {
		sqliteErr := err.(*Error)
		if sqliteErr.Code != ErrConstraint {
			t.Errorf("Wrong basic error code: %d != %d",
				sqliteErr.Code, ErrConstraint)
		}
		if sqliteErr.ExtendedCode != ErrConstraintForeignKey {
			t.Errorf("Wrong extended error code: %d != %d",
				sqliteErr.ExtendedCode, ErrConstraintForeignKey)
		}
	}

}

func TestExtendedErrorCodes_NotNull(t *testing.T) {
	dirName, err := ioutil.TempDir("", "sqlite3-err")
	if err != nil {
		t.Fatal(err)
	}
	defer os.RemoveAll(dirName)

	dbFileName := path.Join(dirName, "test.db")
	db, err := sql.Open("sqlite3", dbFileName)
	if err != nil {
		t.Error(err)
	}
	defer db.Close()

	_, err = db.Exec("PRAGMA foreign_keys=ON;")
	if err != nil {
		t.Errorf("PRAGMA foreign_keys=ON: %v", err)
	}

	_, err = db.Exec(`CREATE TABLE Foo (
		id INTEGER PRIMARY KEY AUTOINCREMENT,
		value INTEGER NOT NULL,
		ref INTEGER NULL REFERENCES Foo (id),
		UNIQUE(value)
	);`)
	if err != nil {
		t.Error(err)
	}

	res, err := db.Exec("INSERT INTO Foo (value) VALUES (100);")
	if err != nil {
		t.Fatalf("Creating first row: %v", err)
	}

	id, err := res.LastInsertId()
	if err != nil {
		t.Fatalf("Retrieving last insert id: %v", err)
	}

	_, err = db.Exec("INSERT INTO Foo (ref) VALUES (?);", id)
	if err == nil {
		t.Error("No error!")
	} else {
		sqliteErr := err.(*Error)
		if sqliteErr.Code != ErrConstraint {
			t.Errorf("Wrong basic error code: %d != %d",
				sqliteErr.Code, ErrConstraint)
		}
		if sqliteErr.ExtendedCode != ErrConstraintNotNull {
			t.Errorf("Wrong extended error code: %d != %d",
				sqliteErr.ExtendedCode, ErrConstraintNotNull)
		}
	}

}

func TestExtendedErrorCodes_Unique(t *testing.T) {
	dirName, err := ioutil.TempDir("", "sqlite3-err")
	if err != nil {
		t.Fatal(err)
	}
	defer os.RemoveAll(dirName)

	dbFileName := path.Join(dirName, "test.db")
	db, err := sql.Open("sqlite3", dbFileName)
	if err != nil {
		t.Error(err)
	}
	defer db.Close()

	_, err = db.Exec("PRAGMA foreign_keys=ON;")
	if err != nil {
		t.Errorf("PRAGMA foreign_keys=ON: %v", err)
	}

	_, err = db.Exec(`CREATE TABLE Foo (
		id INTEGER PRIMARY KEY AUTOINCREMENT,
		value INTEGER NOT NULL,
		ref INTEGER NULL REFERENCES Foo (id),
		UNIQUE(value)
	);`)
	if err != nil {
		t.Error(err)
	}

	res, err := db.Exec("INSERT INTO Foo (value) VALUES (100);")
	if err != nil {
		t.Fatalf("Creating first row: %v", err)
	}

	id, err := res.LastInsertId()
	if err != nil {
		t.Fatalf("Retrieving last insert id: %v", err)
	}

	_, err = db.Exec("INSERT INTO Foo (ref, value) VALUES (?, 100);", id)
	if err == nil {
		t.Error("No error!")
	} else {
		sqliteErr := err.(*Error)
		if sqliteErr.Code != ErrConstraint {
			t.Errorf("Wrong basic error code: %d != %d",
				sqliteErr.Code, ErrConstraint)
		}
		if sqliteErr.ExtendedCode != ErrConstraintUnique {
			t.Errorf("Wrong extended error code: %d != %d",
				sqliteErr.ExtendedCode, ErrConstraintUnique)
		}
		extended := sqliteErr.Code.Extend(3).Error()
		expected := "constraint failed"
		if extended != expected {
			t.Errorf("Wrong basic error code: %q != %q",
				extended, expected)
		}
	}

}
<|endoftext|>


<commit_before>package arn

import (
	"errors"

	"sort"

	"github.com/aerogo/nano"
	"github.com/fatih/color"
)

// Quote ...
type Quote struct {
	ID            string   `json:"id"`
	Description   string   `json:"description" editable:"true" type:"textarea"`
	CharacterID   string   `json:"characterId" editable:"true"`
	AnimeID       string   `json:"animeId" editable:"true"`
	EpisodeNumber int      `json:"episode" editable:"true"`
	Time          int      `json:"time" editable:"true"`
	Likes         []string `json:"likes"`
	IsDraft       bool     `json:"isDraft"`
	Created       string   `json:"created"`
	CreatedBy     string   `json:"createdBy"`
	Edited        string   `json:"edited"`
	EditedBy      string   `json:"editedBy"`
}

// Link returns a single quote.
func (quote *Quote) Link() string {
	return "/quote/" + quote.ID
}

// Creator returns the user who created this quote.
func (quote *Quote) Creator() *User {
	user, _ := GetUser(quote.CreatedBy)
	return user
}

// EditedByUser returns the user who edited this quote.
func (quote *Quote) EditedByUser() *User {
	user, _ := GetUser(quote.EditedBy)
	return user
}

// Publish ...
func (quote *Quote) Publish() error {
	// No draft
	if !quote.IsDraft {
		return errors.New("Not a draft")
	}

	// No description
	if quote.Description == "" {
		return errors.New("A description is required")
	}

	// No character
	if quote.CharacterID == "" {
		return errors.New("A character is required")
	}

	// EpisodeNumber is provided without the anime
	if quote.AnimeID == "" && quote.EpisodeNumber != 0 {
		return errors.New("An anime is required before adding an episode")
	}

	// The time is provided without the EpisodeNumber
	if quote.EpisodeNumber == 0 && quote.Time != 0 {
		return errors.New("An episode is required before adding a time")
	}

	draftIndex, err := GetDraftIndex(quote.CreatedBy)

	if err != nil {
		return err
	}

	if draftIndex.QuoteID == "" {
		return errors.New("Quote draft doesn't exist in the user draft index")
	}

	character, characterErr := GetCharacter(quote.CharacterID)

	if characterErr != nil {
		return errors.New("Character does not exist")
	}

	// Append to quotes Ids
	character.QuoteIDs = append(character.QuoteIDs, quote.ID)

	// Save the character
	character.Save()

	quote.IsDraft = false
	draftIndex.QuoteID = ""
	draftIndex.Save()
	return nil
}

// Unpublish ...
func (quote *Quote) Unpublish() error {
	draftIndex, err := GetDraftIndex(quote.CreatedBy)

	if err != nil {
		return err
	}

	if draftIndex.QuoteID != "" {
		return errors.New("You still have an unfinished draft")
	}

	quote.IsDraft = true
	draftIndex.QuoteID = quote.ID
	draftIndex.Save()
	return nil
}

// Like adds a user to the quote's Likes array if they aren't already in it.
func (quote *Quote) Like(userID string) {
	for _, id := range quote.Likes {
		if id == userID {
			return
		}
	}

	quote.Likes = append(quote.Likes, userID)
}

// Unlike removes the user from the quote's Likes array if they are in it.
func (quote *Quote) Unlike(userID string) {
	for index, id := range quote.Likes {
		if id == userID {
			quote.Likes = append(quote.Likes[:index], quote.Likes[index+1:]...)
			return
		}
	}
}

// LikedBy checks to see if the user has liked the quote.
func (quote *Quote) LikedBy(userID string) bool {
	for _, id := range quote.Likes {
		if id == userID {
			return true
		}
	}

	return false
}

// GetQuote returns a single quote.
func GetQuote(id string) (*Quote, error) {
	obj, err := DB.Get("Quote", id)

	if err != nil {
		return nil, err
	}

	return obj.(*Quote), nil
}

// StreamCompanies returns a stream of all quotes.
func StreamQuotes() chan *Quote {
	channel := make(chan *Quote, nano.ChannelBufferSize)

	go func() {
		for obj := range DB.All("Quote") {
			channel <- obj.(*Quote)
		}

		close(channel)
	}()

	return channel
}

// AllQuotes returns a slice of all quotes.
func AllQuotes() []*Quote {
	var all []*Quote

	stream := StreamQuotes()

	for obj := range stream {
		all = append(all, obj)
	}

	return all
}

// Character returns the character cited in the quote
func (quote *Quote) Character() *Character {
	character, _ := GetCharacter(quote.CharacterID)
	return character
}

// Anime fetches the anime where the quote is said.
func (quote *Quote) Anime() *Anime {
	anime, err := GetAnime(quote.AnimeID)

	if err != nil {
		color.Red("Error fetching anime: %v", err)
	}

	return anime
}

// SortQuotesLatestFirst ...
func SortQuotesLatestFirst(quotes []*Quote) {
	sort.Slice(quotes, func(i, j int) bool {
		return quotes[i].Created > quotes[j].Created
	})
}

// SortQuotesPopularFirst ...
func SortQuotesPopularFirst(quotes []*Quote) {
	sort.Slice(quotes, func(i, j int) bool {
		aLikes := len(quotes[i].Likes)
		bLikes := len(quotes[j].Likes)

		if aLikes == bLikes {
			return quotes[i].Created > quotes[j].Created
		}

		return aLikes > bLikes
	})
}

// FilterQuotes filters all quotes by a custom function.
func FilterQuotes(filter func(*Quote) bool) []*Quote {
	var filtered []*Quote

	for obj := range StreamQuotes() {
		if filter(obj) {
			filtered = append(filtered, obj)
		}
	}

	return filtered
}
<commit_msg>Comment fix<commit_after>package arn

import (
	"errors"

	"sort"

	"github.com/aerogo/nano"
	"github.com/fatih/color"
)

// Quote ...
type Quote struct {
	ID            string   `json:"id"`
	Description   string   `json:"description" editable:"true" type:"textarea"`
	CharacterID   string   `json:"characterId" editable:"true"`
	AnimeID       string   `json:"animeId" editable:"true"`
	EpisodeNumber int      `json:"episode" editable:"true"`
	Time          int      `json:"time" editable:"true"`
	Likes         []string `json:"likes"`
	IsDraft       bool     `json:"isDraft"`
	Created       string   `json:"created"`
	CreatedBy     string   `json:"createdBy"`
	Edited        string   `json:"edited"`
	EditedBy      string   `json:"editedBy"`
}

// Link returns a single quote.
func (quote *Quote) Link() string {
	return "/quote/" + quote.ID
}

// Creator returns the user who created this quote.
func (quote *Quote) Creator() *User {
	user, _ := GetUser(quote.CreatedBy)
	return user
}

// EditedByUser returns the user who edited this quote.
func (quote *Quote) EditedByUser() *User {
	user, _ := GetUser(quote.EditedBy)
	return user
}

// Publish ...
func (quote *Quote) Publish() error {
	// No draft
	if !quote.IsDraft {
		return errors.New("Not a draft")
	}

	// No description
	if quote.Description == "" {
		return errors.New("A description is required")
	}

	// No character
	if quote.CharacterID == "" {
		return errors.New("A character is required")
	}

	// EpisodeNumber is provided without the anime
	if quote.AnimeID == "" && quote.EpisodeNumber != 0 {
		return errors.New("An anime is required before adding an episode")
	}

	// The time is provided without the EpisodeNumber
	if quote.EpisodeNumber == 0 && quote.Time != 0 {
		return errors.New("An episode is required before adding a time")
	}

	draftIndex, err := GetDraftIndex(quote.CreatedBy)

	if err != nil {
		return err
	}

	if draftIndex.QuoteID == "" {
		return errors.New("Quote draft doesn't exist in the user draft index")
	}

	character, characterErr := GetCharacter(quote.CharacterID)

	if characterErr != nil {
		return errors.New("Character does not exist")
	}

	// Append to quotes Ids
	character.QuoteIDs = append(character.QuoteIDs, quote.ID)

	// Save the character
	character.Save()

	quote.IsDraft = false
	draftIndex.QuoteID = ""
	draftIndex.Save()
	return nil
}

// Unpublish ...
func (quote *Quote) Unpublish() error {
	draftIndex, err := GetDraftIndex(quote.CreatedBy)

	if err != nil {
		return err
	}

	if draftIndex.QuoteID != "" {
		return errors.New("You still have an unfinished draft")
	}

	quote.IsDraft = true
	draftIndex.QuoteID = quote.ID
	draftIndex.Save()
	return nil
}

// Like adds a user to the quote's Likes array if they aren't already in it.
func (quote *Quote) Like(userID string) {
	for _, id := range quote.Likes {
		if id == userID {
			return
		}
	}

	quote.Likes = append(quote.Likes, userID)
}

// Unlike removes the user from the quote's Likes array if they are in it.
func (quote *Quote) Unlike(userID string) {
	for index, id := range quote.Likes {
		if id == userID {
			quote.Likes = append(quote.Likes[:index], quote.Likes[index+1:]...)
			return
		}
	}
}

// LikedBy checks to see if the user has liked the quote.
func (quote *Quote) LikedBy(userID string) bool {
	for _, id := range quote.Likes {
		if id == userID {
			return true
		}
	}

	return false
}

// GetQuote returns a single quote.
func GetQuote(id string) (*Quote, error) {
	obj, err := DB.Get("Quote", id)

	if err != nil {
		return nil, err
	}

	return obj.(*Quote), nil
}

// StreamQuotes returns a stream of all quotes.
func StreamQuotes() chan *Quote {
	channel := make(chan *Quote, nano.ChannelBufferSize)

	go func() {
		for obj := range DB.All("Quote") {
			channel <- obj.(*Quote)
		}

		close(channel)
	}()

	return channel
}

// AllQuotes returns a slice of all quotes.
func AllQuotes() []*Quote {
	var all []*Quote

	stream := StreamQuotes()

	for obj := range stream {
		all = append(all, obj)
	}

	return all
}

// Character returns the character cited in the quote
func (quote *Quote) Character() *Character {
	character, _ := GetCharacter(quote.CharacterID)
	return character
}

// Anime fetches the anime where the quote is said.
func (quote *Quote) Anime() *Anime {
	anime, err := GetAnime(quote.AnimeID)

	if err != nil {
		color.Red("Error fetching anime: %v", err)
	}

	return anime
}

// SortQuotesLatestFirst ...
func SortQuotesLatestFirst(quotes []*Quote) {
	sort.Slice(quotes, func(i, j int) bool {
		return quotes[i].Created > quotes[j].Created
	})
}

// SortQuotesPopularFirst ...
func SortQuotesPopularFirst(quotes []*Quote) {
	sort.Slice(quotes, func(i, j int) bool {
		aLikes := len(quotes[i].Likes)
		bLikes := len(quotes[j].Likes)

		if aLikes == bLikes {
			return quotes[i].Created > quotes[j].Created
		}

		return aLikes > bLikes
	})
}

// FilterQuotes filters all quotes by a custom function.
func FilterQuotes(filter func(*Quote) bool) []*Quote {
	var filtered []*Quote

	for obj := range StreamQuotes() {
		if filter(obj) {
			filtered = append(filtered, obj)
		}
	}

	return filtered
}
<|endoftext|>


<commit_before>package mithril

import (
	"fmt"

	"github.com/streadway/amqp"
)

type amqpAdaptedRequest struct {
	Publishing *amqp.Publishing
	Exchange   string
	RoutingKey string
	Mandatory  bool
	Immediate  bool
}

type AMQPHandler struct {
	amqpUri         string
	amqpConn        *amqp.Connection
	handlingChannel *amqp.Channel
	nextHandler     Handler
	confirmAck      chan uint64
	confirmNack     chan uint64
}

func NewAMQPHandler(amqpUri string, next Handler) *AMQPHandler {
	amqpHandler := &AMQPHandler{
		amqpUri: amqpUri,
	}

	amqpHandler.SetNextHandler(next)
	return amqpHandler
}

func (me *AMQPHandler) Init() error {
	var err error

	if err = me.establishConnection(); err != nil {
		return err
	}

	Debugln("AMQP handler initialized")

	if me.nextHandler != nil {
		return me.nextHandler.Init()
	}

	return nil
}

func (me *AMQPHandler) SetNextHandler(handler Handler) {
	me.nextHandler = handler
}

func (me *AMQPHandler) HandleRequest(req *FancyRequest) error {
	var (
		amqpReq *amqpAdaptedRequest
		err     error
	)

	if err = me.establishConnection(); err != nil {
		return err
	}

	amqpReq = me.adaptHttpRequest(req)

	if err = me.publishAdaptedRequest(amqpReq); err != nil {
		Debugln("Failed to publish request:", err)
		return err
	}

	defer me.disconnect()

	return nil
}

func (me *AMQPHandler) establishConnection() error {
	conn, err := amqp.Dial(me.amqpUri)
	if err != nil {
		return err
	}

	me.amqpConn = conn

	handlingChannel, err := me.amqpConn.Channel()
	if err != nil {
		return err
	}

	if err = handlingChannel.Confirm(false); err != nil {
		return err
	}

	me.confirmAck, me.confirmNack = handlingChannel.NotifyConfirm(make(chan uint64, 1), make(chan uint64, 1))

	me.handlingChannel = handlingChannel
	return nil
}

func (me *AMQPHandler) disconnect() {
	if me.handlingChannel != nil {
		me.handlingChannel.Close()
		me.handlingChannel = nil
	}

	if me.amqpConn != nil {
		me.amqpConn.Close()
		me.amqpConn = nil
	}
}

func (me *AMQPHandler) adaptHttpRequest(req *FancyRequest) *amqpAdaptedRequest {
	return &amqpAdaptedRequest{
		Publishing: &amqp.Publishing{
			MessageId:   req.MessageId,
			Timestamp:   req.Timestamp,
			AppId:       req.AppId,
			ContentType: req.ContentType,
			Body:        req.BodyBytes,
		},
		Exchange:   req.Exchange,
		RoutingKey: req.RoutingKey,
		Mandatory:  req.Mandatory,
		Immediate:  req.Immediate,
	}
}

func (me *AMQPHandler) publishAdaptedRequest(amqpReq *amqpAdaptedRequest) error {
	Debugf("Publishing adapted HTTP request %+v\n", amqpReq)

	err := me.handlingChannel.Publish(amqpReq.Exchange,
		amqpReq.RoutingKey, amqpReq.Mandatory,
		amqpReq.Immediate, *amqpReq.Publishing)

	if err != nil {
		return err
	}

	select {
	case _ = <-me.confirmAck:
		return nil
	case _ = <-me.confirmNack:
		return fmt.Errorf("RabbitMQ nack'd message")
	}

	panic("I shouldn't be here")
}
<commit_msg>Moving the deferred disconnect up to where it should be<commit_after>package mithril

import (
	"fmt"

	"github.com/streadway/amqp"
)

type amqpAdaptedRequest struct {
	Publishing *amqp.Publishing
	Exchange   string
	RoutingKey string
	Mandatory  bool
	Immediate  bool
}

type AMQPHandler struct {
	amqpUri         string
	amqpConn        *amqp.Connection
	handlingChannel *amqp.Channel
	nextHandler     Handler
	confirmAck      chan uint64
	confirmNack     chan uint64
}

func NewAMQPHandler(amqpUri string, next Handler) *AMQPHandler {
	amqpHandler := &AMQPHandler{
		amqpUri: amqpUri,
	}

	amqpHandler.SetNextHandler(next)
	return amqpHandler
}

func (me *AMQPHandler) Init() error {
	var err error

	if err = me.establishConnection(); err != nil {
		return err
	}

	Debugln("AMQP handler initialized")

	if me.nextHandler != nil {
		return me.nextHandler.Init()
	}

	return nil
}

func (me *AMQPHandler) SetNextHandler(handler Handler) {
	me.nextHandler = handler
}

func (me *AMQPHandler) HandleRequest(req *FancyRequest) error {
	var (
		amqpReq *amqpAdaptedRequest
		err     error
	)

	if err = me.establishConnection(); err != nil {
		return err
	}

	defer me.disconnect()

	amqpReq = me.adaptHttpRequest(req)

	if err = me.publishAdaptedRequest(amqpReq); err != nil {
		Debugln("Failed to publish request:", err)
		return err
	}

	return nil
}

func (me *AMQPHandler) establishConnection() error {
	conn, err := amqp.Dial(me.amqpUri)
	if err != nil {
		return err
	}

	me.amqpConn = conn

	handlingChannel, err := me.amqpConn.Channel()
	if err != nil {
		return err
	}

	if err = handlingChannel.Confirm(false); err != nil {
		return err
	}

	me.confirmAck, me.confirmNack = handlingChannel.NotifyConfirm(make(chan uint64, 1), make(chan uint64, 1))

	me.handlingChannel = handlingChannel
	return nil
}

func (me *AMQPHandler) disconnect() {
	if me.handlingChannel != nil {
		me.handlingChannel.Close()
		me.handlingChannel = nil
	}

	if me.amqpConn != nil {
		me.amqpConn.Close()
		me.amqpConn = nil
	}
}

func (me *AMQPHandler) adaptHttpRequest(req *FancyRequest) *amqpAdaptedRequest {
	return &amqpAdaptedRequest{
		Publishing: &amqp.Publishing{
			MessageId:   req.MessageId,
			Timestamp:   req.Timestamp,
			AppId:       req.AppId,
			ContentType: req.ContentType,
			Body:        req.BodyBytes,
		},
		Exchange:   req.Exchange,
		RoutingKey: req.RoutingKey,
		Mandatory:  req.Mandatory,
		Immediate:  req.Immediate,
	}
}

func (me *AMQPHandler) publishAdaptedRequest(amqpReq *amqpAdaptedRequest) error {
	Debugf("Publishing adapted HTTP request %+v\n", amqpReq)

	err := me.handlingChannel.Publish(amqpReq.Exchange,
		amqpReq.RoutingKey, amqpReq.Mandatory,
		amqpReq.Immediate, *amqpReq.Publishing)

	if err != nil {
		return err
	}

	select {
	case _ = <-me.confirmAck:
		return nil
	case _ = <-me.confirmNack:
		return fmt.Errorf("RabbitMQ nack'd message")
	}

	panic("I shouldn't be here")
}
<|endoftext|>


<commit_before>package job

import (
	"log"
	"strconv"
	"strings"
	"time"

	"os/exec"

	"os"

	"fmt"

	"github.com/mleonard87/frosty/artifact"
	"github.com/mleonard87/frosty/config"
)

const (
	STATUS_SUCCESS = iota
	STATUS_FAILURE = iota
	BYTES_PER_SI   = 1000
)

var BINARY_SI_UNITS = [...]string{"B", " kB", " MB", " GB", " TB", " PB", " EB", " ZB"}

type JobStatus struct {
	Status            int
	Output            string
	Error             string
	StartTime         time.Time
	EndTime           time.Time
	JobConfig         config.JobConfig
	ArchiveCreated    bool
	ArchiveSize       int64
	TransferStartTime time.Time
	TransferEndTime   time.Time
	TransferError     string
}

func (js JobStatus) ElapsedTime() time.Duration {
	return js.EndTime.Sub(js.StartTime)
}

func (js JobStatus) ElapsedTransferTime() time.Duration {
	return js.TransferEndTime.Sub(js.TransferStartTime)
}

func (js JobStatus) IsSuccessful() bool {
	return js.Status == STATUS_SUCCESS
}

func (js JobStatus) GetArchiveNameDisplay() string {
	return GetArtifactArchiveFileName(js.JobConfig.Name)
}

func (js JobStatus) GetArchiveSizeDisplay() string {
	size := js.ArchiveSize
	for _, unit := range BINARY_SI_UNITS {
		if size < 1024 {
			return strconv.FormatInt(size, 10) + unit
		} else {
			size = size / BYTES_PER_SI
		}
	}
	return strconv.FormatInt(js.ArchiveSize, 10) + BINARY_SI_UNITS[0]
}

func Start(jobConfig config.JobConfig) JobStatus {
	RemoveJobDirectory(jobConfig.Name)
	jobDir, artifactDir := MakeJobDirectories(jobConfig.Name)

	env := os.Environ()
	env = append(env, fmt.Sprintf("FROSTY_JOB_DIR=%s", jobDir))
	env = append(env, fmt.Sprintf("FROSTY_JOB_ARTIFACTS_DIR=%s", artifactDir))
	cmd := exec.Command(jobConfig.Command)
	cmd.Env = env

	js := JobStatus{}

	js.JobConfig = jobConfig
	js.Status = STATUS_SUCCESS
	js.StartTime = time.Now()

	out, err := cmd.Output()
	if err != nil {
		js.Status = STATUS_FAILURE
		js.Error = strings.TrimSpace(err.Error())
	}

	js.EndTime = time.Now()
	js.Output = strings.TrimSpace(string(out[:]))

	archiveTarget := GetArtifactArchiveTargetName(jobConfig.Name)
	js.ArchiveCreated = artifact.MakeArtifactArchive(artifactDir, archiveTarget)

	if js.ArchiveCreated {
		fileInfo, err := os.Stat(archiveTarget)
		if err != nil {
			log.Fatal(err)
		}
		js.ArchiveSize = fileInfo.Size()
	}

	return js
}
<commit_msg>Capture everything written to stderr as well as just the go error message so that this can be reported on. This might be confusing as stdout and stderr are capture separately and displayed separately in the email report so it might not be in order.<commit_after>package job

import (
	"log"
	"strconv"
	"strings"
	"time"

	"os/exec"

	"os"

	"fmt"

	"github.com/mleonard87/frosty/artifact"
	"github.com/mleonard87/frosty/config"
)

const (
	STATUS_SUCCESS = iota
	STATUS_FAILURE = iota
	BYTES_PER_SI   = 1000
)

var BINARY_SI_UNITS = [...]string{"B", " kB", " MB", " GB", " TB", " PB", " EB", " ZB"}

type JobStatus struct {
	Status            int
	Output            string
	Error             string
	StartTime         time.Time
	EndTime           time.Time
	JobConfig         config.JobConfig
	ArchiveCreated    bool
	ArchiveSize       int64
	TransferStartTime time.Time
	TransferEndTime   time.Time
	TransferError     string
}

func (js JobStatus) ElapsedTime() time.Duration {
	return js.EndTime.Sub(js.StartTime)
}

func (js JobStatus) ElapsedTransferTime() time.Duration {
	return js.TransferEndTime.Sub(js.TransferStartTime)
}

func (js JobStatus) IsSuccessful() bool {
	return js.Status == STATUS_SUCCESS
}

func (js JobStatus) GetArchiveNameDisplay() string {
	return GetArtifactArchiveFileName(js.JobConfig.Name)
}

func (js JobStatus) GetArchiveSizeDisplay() string {
	size := js.ArchiveSize
	for _, unit := range BINARY_SI_UNITS {
		if size < 1024 {
			return strconv.FormatInt(size, 10) + unit
		} else {
			size = size / BYTES_PER_SI
		}
	}
	return strconv.FormatInt(js.ArchiveSize, 10) + BINARY_SI_UNITS[0]
}

func Start(jobConfig config.JobConfig) JobStatus {
	RemoveJobDirectory(jobConfig.Name)
	jobDir, artifactDir := MakeJobDirectories(jobConfig.Name)

	env := os.Environ()
	env = append(env, fmt.Sprintf("FROSTY_JOB_DIR=%s", jobDir))
	env = append(env, fmt.Sprintf("FROSTY_JOB_ARTIFACTS_DIR=%s", artifactDir))
	cmd := exec.Command(jobConfig.Command)
	cmd.Env = env

	js := JobStatus{}

	js.JobConfig = jobConfig
	js.Status = STATUS_SUCCESS
	js.StartTime = time.Now()

	out, err := cmd.Output()
	if err != nil {
		js.Status = STATUS_FAILURE

		if ee, ok := err.(*exec.ExitError); ok {
			// Capture and trim any errors logged to stderr.
			se := strings.TrimSpace(string(ee.Stderr))
			// Report both the Go error message and that from stderr
			em := fmt.Sprintf("%s\n%s", se, err.Error())
			js.Error = em
		}
	}

	js.EndTime = time.Now()
	js.Output = strings.TrimSpace(string(out[:]))

	archiveTarget := GetArtifactArchiveTargetName(jobConfig.Name)
	js.ArchiveCreated = artifact.MakeArtifactArchive(artifactDir, archiveTarget)

	if js.ArchiveCreated {
		fileInfo, err := os.Stat(archiveTarget)
		if err != nil {
			log.Fatal(err)
		}
		js.ArchiveSize = fileInfo.Size()
	}

	return js
}
<|endoftext|>


<commit_before>package joystick

import "errors"

func newOsJoystick() osJoystick {
	return osJoystick{}
}

type osJoystick struct {
}

func osinit() {}

func (j *Joystick) prepare() error {
	return errors.New("OS not supported")
}

func (j *Joystick) getState() (*State, error) {
	return nil, errors.New("OS not supported")
}

func (j *Joystick) vibrate(left, right uint16) error {
	return errors.New("OS not supported")
}

func (j *Joystick) close() error {
	return errors.New("OS not supported")
}

func getJoysticks() []*Joystick {
	return nil
}
<commit_msg>Add libudev implementation for linux<commit_after>package joystick

import (
	"github.com/oakmound/oak/dlog"
	"errors"
	"os"

	"encoding/binary"
	"fmt"
	"path/filepath"
	"regexp"
	"runtime"

	"github.com/citilinkru/libudev"
	"github.com/citilinkru/libudev/types"
)

// This has all been tested with wired xbox 360 controllers.
// Todo: get more controllers, test with more controllers.

func newJoystick(devName string, id uint32) *Joystick {
	return &Joystick{
		Handler:  event.DefaultBus,
		PollRate: timing.FPSToDuration(60),
		id:       id,
		osJoystick: osJoystick{
			devName: devName,
			quit: make(chan struct{})
		},
	}
}
type osJoystick struct {
	devName string
	fh *os.File
	cache State
	quit chan struct{}
}

func osinit() {}

type event struct {
	Time   uint32
	Value  int16
	Type   uint8
	Number uint8
}

const (
	axisType = 2
	buttonType = 1
)

var (
	buttons []string{
		0: "A",
		1: "B",
		2: "X",
		3: "Y",
		4: "LeftShoulder",
		5: "RightShoulder",
		6: "Back",
		7: "Start",
		// 8 is the "Xbox" button in the center
		9: "LeftStick",
		10: "RightStick",
	}
)

func (j *Joystick) prepare() error {
	var err error
	j.fh, err = os.Open(devName)
	if err == nil {
		go func(j *Joystick) {
			// Read events continually
			e := &event{}
			for {
				select {
				case <-j.quit:
						return
				}
				err := binary.Read(j.fh, binary.LittleEndian, &e)
				dlog.ErrorCheck(err)
				switch e.Type {
					case axis:
						switch e.Number {
						case 0:
							j.cache.StickLX = e.Value
						case 1:
							j.cache.StickLY = e.Value
						case 2:
							// The controller offers int16 fidelity of the
							// triggers. We're lowering it to Xinput's uint8
							// Todo: Flip that around?
							j.cache.TriggerL = uint8(uint16(e.Value)/16)
						case 3:
							j.cache.StickRX = e.Value
						case 4:
							j.cache.StickRY = e.Value
						case 5:
							j.cache.TriggerR = uint8(uint16(e.Value)/16)
						case 6:
							if e.Value < 0 {
								j.cache.Buttons["Left"] = true
								j.cache.Buttons["Right"] = false
							} else if e.Value > 0 {
								j.cache.Buttons["Right"] = true
								j.cache.Buttons["Left"] = false
							} else {
								j.cache.Buttons["Right"] = false
								j.cache.Buttons["Left"] = false
							}
						case 7:
							if e.Value < 0 {
								j.cache.Buttons["Up"] = true
								j.cache.Buttons["Down"] = false
							} else if e.Value > 0 {
								j.cache.Buttons["Down"] = true
								j.cache.Buttons["Up"] = false
							} else {
								j.cache.Buttons["Down"] = false
								j.cache.Buttons["Up"] = false
							}
						}
					case button:
						j.cache.Buttons[buttons[e.Number]] = (e.Value == 1)
				}
				// No mutex here could cause a frame delay on inputs
				j.cache.Frame = e.Time
			}
		}(j)
	}
	return err
}

func (j *Joystick) getState() (*State, error) {
	s := new(State)
	*s = j.cache
	return s, nil
}

func (j *Joystick) vibrate(left, right uint16) error {
	return errors.New("Vibration not supported")
}

func (j *Joystick) close() error {
	go func() {
		j.quit <- struct{}{}
	}
	return j.fh.Close()
}

func getJoysticks() []*Joystick {
	sc := libudev.NewScanner()
	err, dvs := sc.ScanDevices()
	if err != nil {
		fmt.Println(err)
		return
	}
	// Joysticks contain "js%d"
	rgx, err := regexp.Compile("js(\\d)+")
	if err != nil {
		dlog.Error(err)
		return nil
	}

	filtered := []*types.Device{}

	for _, d := range dvs {
		// Find joysticks
		if !rgx.MatchString(d.Devpath) {
			continue
		}
		// Ignore mice
		if v, ok := d.Env["ID_INPUT_MOUSE"]; ok && v == "1" {
			continue
		}
		// Todo: what else do we ignore?
		fmt.Printf("%+v\n", d)
		filtered = append(filtered, d)
	}

	joys := make([]*Joystick, len(filtered))
	for i, f := range filtered {
		var id uint32 = math.MaxUint32
		matches := rgx.FindStringSubmatch(d.Devpath)
		if len(matches) > 1 {
			id, err = strconv.Atoi(matches[1])
			dlog.ErrorCheck(err)
		}
		joys[i] = newJoystick(filepath.Join("/", "dev", d.Env["DEVNAME"]), id)
	}
	return joys
}<|endoftext|>


<commit_before>package henchman

import (
	"bytes"
	"encoding/json"
	"fmt"
	log "gopkg.in/Sirupsen/logrus.v0"
	"io/ioutil"
	"os"
	"os/exec"
	"path"
	"strconv"

	"github.com/flosch/pongo2"
	"github.com/pborman/uuid"
)

type Task struct {
	Id           string
	Sudo         bool
	Name         string
	Module       *Module
	IgnoreErrors bool `yaml:"ignore_errors"`
	Local        bool
	When         string
	Register     string
	Vars         VarsMap
}

type TaskResult struct {
	State  string      `json:"status"`
	Output interface{} `json:"output,omitempty"`
	Msg    string      `json:"msg"`
}

func getTaskResult(buf *bytes.Buffer) (*TaskResult, error) {
	var taskResult TaskResult
	resultInBytes := []byte(buf.String())
	err := json.Unmarshal(resultInBytes, &taskResult)
	if err != nil {
		return &TaskResult{}, err
	}
	return &taskResult, nil
}

func setTaskResult(taskResult *TaskResult, buf *bytes.Buffer) error {
	resultInBytes := []byte(buf.String())
	err := json.Unmarshal(resultInBytes, &taskResult)
	if err != nil {
		return err
	}
	return nil
}

// strings will be evaluated using pongo2 templating with context of
// VarsMap and RegisterMap
func renderValue(value string, varsMap VarsMap, registerMap map[string]interface{}) (string, error) {
	tmpl, err := pongo2.FromString(value)
	if err != nil {
		return "", fmt.Errorf("Templating :: %s", err.Error())
	}

	ctxt := pongo2.Context{"vars": varsMap}
	ctxt = ctxt.Update(registerMap)

	out, err := tmpl.Execute(ctxt)
	if err != nil {
		return "", fmt.Errorf("Executing :: %s", err.Error())
	}
	return out, nil
}

// wrapper for Rendering each task
func (task *Task) Render(vars VarsMap, registerMap RegMap) error {
	var err error
	task.Name, err = renderValue(task.Name, vars, registerMap)
	if err != nil {
		return err
	}

	if task.When != "" {
		task.When, err = renderValue("{{"+task.When+"}}", vars, registerMap)
		if err != nil {
			return err
		}
	}

	for k, v := range task.Module.Params {
		task.Module.Params[k], err = renderValue(v, vars, registerMap)
		if err != nil {
			return err
		}
	}

	return nil
}

// checks and converts when to bool
func (task *Task) ProcessWhen() (bool, error) {
	if task.When == "" {
		return true, nil
	}

	result, err := strconv.ParseBool(task.When)
	if err != nil {
		return false, err
	}

	return result, nil
}

func (task *Task) Run(machine *Machine, vars VarsMap, registerMap RegMap) (*TaskResult, error) {
	// Add current host to vars
	task.Id = uuid.New()

	//local copy of module params for each machine
	moduleParams := make(map[string]string)
	for k, v := range task.Module.Params {
		moduleParams[k] = v
	}
	proceed, err := task.ProcessWhen()
	if err != nil {
		return &TaskResult{}, err
	}

	if proceed == false {
		return &TaskResult{State: "skipped"}, nil
	}

	modPath, _, err := task.Module.Resolve()
	if err != nil {
		return &TaskResult{}, fmt.Errorf("Module Path :: %s", err.Error())
	}

	execOrder, err := task.Module.ExecOrder()
	// currently err will always be nil
	if err != nil {
		return &TaskResult{}, fmt.Errorf("Exec Order :: %s", err.Error())
	}

	remoteModDir := "${HOME}/.henchman/"
	remoteModPath := path.Join(remoteModDir, task.Module.Name)
	// NOTE: Info or Debug level
	Debug(log.Fields{
		"task":   task.Name,
		"host":   task.Vars["current_host"],
		"module": task.Module.Name,
		"order":  execOrder,
	}, "Exec Order")

	var taskResult TaskResult
	for _, execStep := range execOrder {
		switch execStep {
		// Exec Order for Default
		case "create_dir":
			// creates remote .henchman location
			_, err = machine.Transport.Exec(fmt.Sprintf("mkdir -p %s", remoteModDir),
				nil, false)
			if err != nil {
				return &TaskResult{}, fmt.Errorf("Creating Mod Path :: %s", err.Error())
			}

		case "put_module":
			// copies module from local location to remote location
			err = machine.Transport.Put(modPath, remoteModDir, "dir")
			if err != nil {
				return &TaskResult{}, fmt.Errorf("Putting Module :: %s", err.Error())
			}

		case "exec_module":
			// executes module by calling the copied module remotely
			log.WithFields(log.Fields{
				"mod path": remoteModPath,
				"hosts":    task.Vars["current_host"],
				"task":     task.Name,
				"module":   task.Module.Name,
			}).Info("Executing Module in Task")

			jsonParams, err := json.Marshal(moduleParams)
			if err != nil {
				return &TaskResult{}, fmt.Errorf("Exec Module :: Json :: %s", err.Error())
			}
			buf, err := machine.Transport.Exec(remoteModPath, jsonParams, task.Sudo)
			if err != nil {
				return &TaskResult{}, fmt.Errorf("Exec Module :: %s", err.Error())
			}
			//This should not be empty
			err = setTaskResult(&taskResult, buf)
			if err != nil {
				return &TaskResult{}, err
			}

			// Exec Order for dependencies
		case "tar_module":
			// check if the local tar module exists
			// else create the local tar of the module
			if _, err := os.Stat(modPath + ".tar"); os.IsNotExist(err) {
				curDir, err := os.Getwd()
				if err != nil {
					return &TaskResult{}, fmt.Errorf("Tarring Module :: GetWd :: %s", err.Error())
				}

				os.Chdir(modPath)
				os.Chdir("../")
				cmd := "tar"
				args := []string{"-cvf", task.Module.Name + ".tar", task.Module.Name}
				if err := exec.Command(cmd, args...).Run(); err != nil {
					return &TaskResult{}, fmt.Errorf("Tarring Module :: %s", err.Error())
				}
				os.Chdir(curDir)
			}
		case "put_tar_module":
			// copies module from local location to remote location
			err = machine.Transport.Put(modPath+".tar", remoteModDir, "dir")
			if err != nil {
				return &TaskResult{}, fmt.Errorf("Putting Tar Module :: %s", err.Error())
			}

			// deletes module.tar from local modules folder
			/*
				cmd := "rm"
				args := []string{modPath + ".tar"}
				if err := exec.Command(cmd, args...).Run(); err != nil {
					return &TaskResult{}, fmt.Errorf("Putting Tar Module :: %s", err.Error())
				}
			*/
		case "untar_module":
			// untars the module
			cmd := fmt.Sprintf("tar -xvf %s -C %s", remoteModPath+".tar", remoteModDir)
			_, err := machine.Transport.Exec(cmd, nil, task.Sudo)
			if err != nil {
				return &TaskResult{}, fmt.Errorf("Untar Module :: %s", err.Error())
			}

			cmd = fmt.Sprintf("/bin/rm %s", remoteModPath+".tar")
			_, err = machine.Transport.Exec(cmd, nil, task.Sudo)
			if err != nil {
				return &TaskResult{}, fmt.Errorf("Untar Module :: %s", err.Error())
			}
		case "exec_tar_module":
			// executes module by calling the copied module remotely
			// NOTE: may want to just change the way remoteModPath is created
			newModPath := remoteModDir + task.Module.Name + "/exec"
			log.WithFields(log.Fields{
				"mod path": newModPath,
				"host":     task.Vars["current_host"],
				"task":     task.Name,
				"module":   task.Module.Name,
			}).Info("Executing Module in Task")

			jsonParams, err := json.Marshal(moduleParams)
			if err != nil {
				return &TaskResult{}, fmt.Errorf("Exec Tar Module :: Json :: %s", err.Error())
			}
			buf, err := machine.Transport.Exec(newModPath, jsonParams, task.Sudo)
			if err != nil {
				return &TaskResult{}, fmt.Errorf("Exec Tar Module :: %s", err.Error())
			}

			//This should not be empty
			err = setTaskResult(&taskResult, buf)
			if err != nil {
				return &TaskResult{}, fmt.Errorf("Exec Tar Module :: %s", err.Error())
			}
		case "copy_remote":
			//copies file from remote .henchman location to expected location
			remoteSrcPath, present := moduleParams["src"]
			if !present {
				return &TaskResult{}, fmt.Errorf("Unable to find 'src' parameter")
			}

			dstPath, present := moduleParams["dest"]
			if !present {
				return &TaskResult{}, fmt.Errorf("Unable to find 'dest' parameter")
			}

			_, localSrcFile := path.Split(remoteSrcPath)
			srcPath := path.Join(remoteModDir, localSrcFile)

			cmd := fmt.Sprintf("/bin/cp %s %s", srcPath, dstPath)
			buf, err := machine.Transport.Exec(cmd, nil, task.Sudo)
			if err != nil {
				return &TaskResult{}, err
			}
			if len(buf.String()) != 0 {
				err = setTaskResult(&taskResult, buf)
				if err != nil {
					return &TaskResult{}, err
				}
				//log.Println(taskResult)
			}
		case "put_file":
			//scp's file from local location to remote location
			srcPath, present := moduleParams["src"]
			if !present {
				return &TaskResult{}, fmt.Errorf("Unable to find 'src' parameter")
			}
			_, srcFile := path.Split(srcPath)
			dstPath := path.Join(remoteModDir, srcFile)

			err = machine.Transport.Put(srcPath, dstPath, "file")

			if err != nil {
				return &TaskResult{}, fmt.Errorf("Putting File :: %s", err.Error())
			}

		case "process_template":
			srcPath, present := moduleParams["src"]
			if !present {
				return &TaskResult{}, fmt.Errorf("Unable to find 'src' parameter")
			}

			tpl, err := pongo2.FromFile(srcPath)
			if err != nil {
				return &TaskResult{}, fmt.Errorf("Process Template :: %s", err.Error())
			}
			out, err := tpl.Execute(pongo2.Context{"vars": vars})
			if err != nil {
				return &TaskResult{}, fmt.Errorf("Process Template :: %s", err.Error())
			}
			tmpDir, srcFile := path.Split(srcPath)
			srcFile = srcFile + "_" + machine.Hostname
			tmpFileName := fmt.Sprintf(".%s", srcFile)
			tmpFile := path.Join(tmpDir, tmpFileName)

			err = ioutil.WriteFile(tmpFile, []byte(out), 0644)
			if err != nil {
				return &TaskResult{}, fmt.Errorf("Process Template :: %s", err.Error())
			}
			moduleParams["srcOrig"] = srcPath
			moduleParams["src"] = tmpFile
		case "reset_src":
			moduleParams["src"] = moduleParams["srcOrig"]
		}
	}

	// Set to status to ignored if the result is a failure
	if task.IgnoreErrors && (taskResult.State != "ok") {
		taskResult.State = "ignored"
	}

	if task.Register != "" {
		registerMap[task.Register] = taskResult
	}

	return &taskResult, nil
}
<commit_msg>fixed mismatch error<commit_after>package henchman

import (
	"bytes"
	"encoding/json"
	"fmt"
	log "gopkg.in/Sirupsen/logrus.v0"
	"io/ioutil"
	"os"
	"os/exec"
	"path"
	"strconv"

	"github.com/flosch/pongo2"
	"github.com/pborman/uuid"
)

type Task struct {
	Id           string
	Sudo         bool
	Name         string
	Module       *Module
	IgnoreErrors bool `yaml:"ignore_errors"`
	Local        bool
	When         string
	Register     string
	Vars         VarsMap
}

type TaskResult struct {
	State  string      `json:"status"`
	Output interface{} `json:"output,omitempty"`
	Msg    string      `json:"msg"`
}

func getTaskResult(buf *bytes.Buffer) (*TaskResult, error) {
	var taskResult TaskResult
	resultInBytes := []byte(buf.String())
	err := json.Unmarshal(resultInBytes, &taskResult)
	if err != nil {
		return &TaskResult{}, err
	}
	return &taskResult, nil
}

func setTaskResult(taskResult *TaskResult, buf *bytes.Buffer) error {
	resultInBytes := []byte(buf.String())
	err := json.Unmarshal(resultInBytes, &taskResult)
	if err != nil {
		return err
	}
	return nil
}

// strings will be evaluated using pongo2 templating with context of
// VarsMap and RegisterMap
func renderValue(value string, varsMap VarsMap, registerMap map[string]interface{}) (string, error) {
	tmpl, err := pongo2.FromString(value)
	if err != nil {
		return "", fmt.Errorf("Templating :: %s", err.Error())
	}

	ctxt := pongo2.Context{"vars": varsMap}
	ctxt = ctxt.Update(registerMap)

	out, err := tmpl.Execute(ctxt)
	if err != nil {
		return "", fmt.Errorf("Executing :: %s", err.Error())
	}
	return out, nil
}

// wrapper for Rendering each task
func (task *Task) Render(vars VarsMap, registerMap RegMap) error {
	var err error
	task.Name, err = renderValue(task.Name, vars, registerMap)
	if err != nil {
		return err
	}

	if task.When != "" {
		task.When, err = renderValue("{{"+task.When+"}}", vars, registerMap)
		if err != nil {
			return err
		}
	}

	for k, v := range task.Module.Params {
		task.Module.Params[k], err = renderValue(v, vars, registerMap)
		if err != nil {
			return err
		}
	}

	return nil
}

// checks and converts when to bool
func (task *Task) ProcessWhen() (bool, error) {
	if task.When == "" {
		return true, nil
	}

	result, err := strconv.ParseBool(task.When)
	if err != nil {
		return false, err
	}

	return result, nil
}

func (task *Task) Run(machine *Machine, vars VarsMap, registerMap RegMap) (*TaskResult, error) {
	// Add current host to vars
	task.Id = uuid.New()

	//local copy of module params for each machine
	moduleParams := make(map[string]string)
	for k, v := range task.Module.Params {
		moduleParams[k] = v
	}
	proceed, err := task.ProcessWhen()
	if err != nil {
		return &TaskResult{}, err
	}

	if proceed == false {
		return &TaskResult{State: "skipped"}, nil
	}

	modPath, err := task.Module.Resolve()
	if err != nil {
		return &TaskResult{}, fmt.Errorf("Module Path :: %s", err.Error())
	}

	execOrder, err := task.Module.ExecOrder()
	// currently err will always be nil
	if err != nil {
		return &TaskResult{}, fmt.Errorf("Exec Order :: %s", err.Error())
	}

	remoteModDir := "${HOME}/.henchman/"
	remoteModPath := path.Join(remoteModDir, task.Module.Name)
	// NOTE: Info or Debug level
	Debug(log.Fields{
		"task":   task.Name,
		"host":   task.Vars["current_host"],
		"module": task.Module.Name,
		"order":  execOrder,
	}, "Exec Order")

	var taskResult TaskResult
	for _, execStep := range execOrder {
		switch execStep {
		// Exec Order for Default
		case "create_dir":
			// creates remote .henchman location
			_, err = machine.Transport.Exec(fmt.Sprintf("mkdir -p %s", remoteModDir),
				nil, false)
			if err != nil {
				return &TaskResult{}, fmt.Errorf("Creating Mod Path :: %s", err.Error())
			}

		case "put_module":
			// copies module from local location to remote location
			err = machine.Transport.Put(modPath, remoteModDir, "dir")
			if err != nil {
				return &TaskResult{}, fmt.Errorf("Putting Module :: %s", err.Error())
			}

		case "exec_module":
			// executes module by calling the copied module remotely
			log.WithFields(log.Fields{
				"mod path": remoteModPath,
				"hosts":    task.Vars["current_host"],
				"task":     task.Name,
				"module":   task.Module.Name,
			}).Info("Executing Module in Task")

			jsonParams, err := json.Marshal(moduleParams)
			if err != nil {
				return &TaskResult{}, fmt.Errorf("Exec Module :: Json :: %s", err.Error())
			}
			buf, err := machine.Transport.Exec(remoteModPath, jsonParams, task.Sudo)
			if err != nil {
				return &TaskResult{}, fmt.Errorf("Exec Module :: %s", err.Error())
			}
			//This should not be empty
			err = setTaskResult(&taskResult, buf)
			if err != nil {
				return &TaskResult{}, err
			}

			// Exec Order for dependencies
		case "tar_module":
			// check if the local tar module exists
			// else create the local tar of the module
			if _, err := os.Stat(modPath + ".tar"); os.IsNotExist(err) {
				curDir, err := os.Getwd()
				if err != nil {
					return &TaskResult{}, fmt.Errorf("Tarring Module :: GetWd :: %s", err.Error())
				}

				os.Chdir(modPath)
				os.Chdir("../")
				cmd := "tar"
				args := []string{"-cvf", task.Module.Name + ".tar", task.Module.Name}
				if err := exec.Command(cmd, args...).Run(); err != nil {
					return &TaskResult{}, fmt.Errorf("Tarring Module :: %s", err.Error())
				}
				os.Chdir(curDir)
			}
		case "put_tar_module":
			// copies module from local location to remote location
			err = machine.Transport.Put(modPath+".tar", remoteModDir, "dir")
			if err != nil {
				return &TaskResult{}, fmt.Errorf("Putting Tar Module :: %s", err.Error())
			}

			// deletes module.tar from local modules folder
			/*
				cmd := "rm"
				args := []string{modPath + ".tar"}
				if err := exec.Command(cmd, args...).Run(); err != nil {
					return &TaskResult{}, fmt.Errorf("Putting Tar Module :: %s", err.Error())
				}
			*/
		case "untar_module":
			// untars the module
			cmd := fmt.Sprintf("tar -xvf %s -C %s", remoteModPath+".tar", remoteModDir)
			_, err := machine.Transport.Exec(cmd, nil, task.Sudo)
			if err != nil {
				return &TaskResult{}, fmt.Errorf("Untar Module :: %s", err.Error())
			}

			cmd = fmt.Sprintf("/bin/rm %s", remoteModPath+".tar")
			_, err = machine.Transport.Exec(cmd, nil, task.Sudo)
			if err != nil {
				return &TaskResult{}, fmt.Errorf("Untar Module :: %s", err.Error())
			}
		case "exec_tar_module":
			// executes module by calling the copied module remotely
			// NOTE: may want to just change the way remoteModPath is created
			newModPath := remoteModDir + task.Module.Name + "/exec"
			log.WithFields(log.Fields{
				"mod path": newModPath,
				"host":     task.Vars["current_host"],
				"task":     task.Name,
				"module":   task.Module.Name,
			}).Info("Executing Module in Task")

			jsonParams, err := json.Marshal(moduleParams)
			if err != nil {
				return &TaskResult{}, fmt.Errorf("Exec Tar Module :: Json :: %s", err.Error())
			}
			buf, err := machine.Transport.Exec(newModPath, jsonParams, task.Sudo)
			if err != nil {
				return &TaskResult{}, fmt.Errorf("Exec Tar Module :: %s", err.Error())
			}

			//This should not be empty
			err = setTaskResult(&taskResult, buf)
			if err != nil {
				return &TaskResult{}, fmt.Errorf("Exec Tar Module :: %s", err.Error())
			}
		case "copy_remote":
			//copies file from remote .henchman location to expected location
			remoteSrcPath, present := moduleParams["src"]
			if !present {
				return &TaskResult{}, fmt.Errorf("Unable to find 'src' parameter")
			}

			dstPath, present := moduleParams["dest"]
			if !present {
				return &TaskResult{}, fmt.Errorf("Unable to find 'dest' parameter")
			}

			_, localSrcFile := path.Split(remoteSrcPath)
			srcPath := path.Join(remoteModDir, localSrcFile)

			cmd := fmt.Sprintf("/bin/cp %s %s", srcPath, dstPath)
			buf, err := machine.Transport.Exec(cmd, nil, task.Sudo)
			if err != nil {
				return &TaskResult{}, err
			}
			if len(buf.String()) != 0 {
				err = setTaskResult(&taskResult, buf)
				if err != nil {
					return &TaskResult{}, err
				}
				//log.Println(taskResult)
			}
		case "put_file":
			//scp's file from local location to remote location
			srcPath, present := moduleParams["src"]
			if !present {
				return &TaskResult{}, fmt.Errorf("Unable to find 'src' parameter")
			}
			_, srcFile := path.Split(srcPath)
			dstPath := path.Join(remoteModDir, srcFile)

			err = machine.Transport.Put(srcPath, dstPath, "file")

			if err != nil {
				return &TaskResult{}, fmt.Errorf("Putting File :: %s", err.Error())
			}

		case "process_template":
			srcPath, present := moduleParams["src"]
			if !present {
				return &TaskResult{}, fmt.Errorf("Unable to find 'src' parameter")
			}

			tpl, err := pongo2.FromFile(srcPath)
			if err != nil {
				return &TaskResult{}, fmt.Errorf("Process Template :: %s", err.Error())
			}
			out, err := tpl.Execute(pongo2.Context{"vars": vars})
			if err != nil {
				return &TaskResult{}, fmt.Errorf("Process Template :: %s", err.Error())
			}
			tmpDir, srcFile := path.Split(srcPath)
			srcFile = srcFile + "_" + machine.Hostname
			tmpFileName := fmt.Sprintf(".%s", srcFile)
			tmpFile := path.Join(tmpDir, tmpFileName)

			err = ioutil.WriteFile(tmpFile, []byte(out), 0644)
			if err != nil {
				return &TaskResult{}, fmt.Errorf("Process Template :: %s", err.Error())
			}
			moduleParams["srcOrig"] = srcPath
			moduleParams["src"] = tmpFile
		case "reset_src":
			moduleParams["src"] = moduleParams["srcOrig"]
		}
	}

	// Set to status to ignored if the result is a failure
	if task.IgnoreErrors && (taskResult.State != "ok") {
		taskResult.State = "ignored"
	}

	if task.Register != "" {
		registerMap[task.Register] = taskResult
	}

	return &taskResult, nil
}
<|endoftext|>


<commit_before>package gin

import (
	"bufio"
	"errors"
	"log"
	"net"
	"net/http"
)

type (
	ResponseWriter interface {
		http.ResponseWriter
		Status() int
		Written() bool
		WriteHeaderNow()
		Hijack() (net.Conn, *bufio.ReadWriter, error)
	}

	responseWriter struct {
		http.ResponseWriter
		status  int
		written bool
	}
)

func (w *responseWriter) reset(writer http.ResponseWriter) {
	w.ResponseWriter = writer
	w.status = 200
	w.written = false
}

func (w *responseWriter) WriteHeader(code int) {
	if code > 0 {
		w.status = code
		if w.written {
			log.Println("[GIN] WARNING. Headers were already written!")
		}
	}
}

func (w *responseWriter) WriteHeaderNow() {
	if !w.written {
		w.written = true
		w.ResponseWriter.WriteHeader(w.status)
	}
}

func (w *responseWriter) Write(data []byte) (n int, err error) {
	w.WriteHeaderNow()
	return w.ResponseWriter.Write(data)
}

func (w *responseWriter) Status() int {
	return w.status
}

func (w *responseWriter) Written() bool {
	return w.written
}

// allow connection hijacking
func (w *responseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
	hijacker, ok := w.ResponseWriter.(http.Hijacker)
	if !ok {
		return nil, nil, errors.New("the ResponseWriter doesn't support the Hijacker interface")
	}
	return hijacker.Hijack()
}
<commit_msg>Referencing the http.Hijacker interface<commit_after>package gin

import (
	"bufio"
	"errors"
	"log"
	"net"
	"net/http"
)

type (
	ResponseWriter interface {
		http.ResponseWriter
		http.Hijacker
		Status() int
		Written() bool
		WriteHeaderNow()
	}

	responseWriter struct {
		http.ResponseWriter
		status  int
		written bool
	}
)

func (w *responseWriter) reset(writer http.ResponseWriter) {
	w.ResponseWriter = writer
	w.status = 200
	w.written = false
}

func (w *responseWriter) WriteHeader(code int) {
	if code > 0 {
		w.status = code
		if w.written {
			log.Println("[GIN] WARNING. Headers were already written!")
		}
	}
}

func (w *responseWriter) WriteHeaderNow() {
	if !w.written {
		w.written = true
		w.ResponseWriter.WriteHeader(w.status)
	}
}

func (w *responseWriter) Write(data []byte) (n int, err error) {
	w.WriteHeaderNow()
	return w.ResponseWriter.Write(data)
}

func (w *responseWriter) Status() int {
	return w.status
}

func (w *responseWriter) Written() bool {
	return w.written
}

// Implements the http.Hijacker interface
func (w *responseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
	hijacker, ok := w.ResponseWriter.(http.Hijacker)
	if !ok {
		return nil, nil, errors.New("the ResponseWriter doesn't support the Hijacker interface")
	}
	return hijacker.Hijack()
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/mndrix/tap-go"
	"github.com/opencontainers/runtime-tools/cgroups"
	"github.com/opencontainers/runtime-tools/validation/util"
)

func main() {
	var id, prio uint32 = 255, 10
	ifName := "lo"

	t := tap.New()
	t.Header(0)
	defer t.AutoPlan()

	g, err := util.GetDefaultGenerator()
	if err != nil {
		util.Fatal(err)
	}
	g.SetLinuxCgroupsPath(cgroups.AbsCgroupPath)
	g.SetLinuxResourcesNetworkClassID(id)
	g.AddLinuxResourcesNetworkPriorities(ifName, prio)
	err = util.RuntimeOutsideValidate(g, t, util.ValidateLinuxResourcesNetwork)
	if err != nil {
		util.Fatal(err)
	}
}
<commit_msg>validation: add more test cases for linux_cgroups_network<commit_after>package main

import (
	"fmt"
	"net"
	"runtime"

	"github.com/mndrix/tap-go"
	"github.com/opencontainers/runtime-tools/cgroups"
	"github.com/opencontainers/runtime-tools/validation/util"
)

func testNetworkCgroups() error {
	t := tap.New()
	t.Header(0)
	defer t.AutoPlan()

	loIfName := ""
	ethIfName := ""

	// we try to get the first 2 interfaces on the system, to assign
	// the first one to loIfName, the second to ethIfName.
	l, err := net.Interfaces()
	if err != nil {
		// fall back to the default list only with lo
		loIfName = "lo"
		ethIfName = "eth0"
	} else {
		loIfName = l[0].Name
		ethIfName = l[1].Name
	}

	cases := []struct {
		classid    uint32
		prio       uint32
		ifName     string
		withNetNs  bool
		withUserNs bool
	}{
		{255, 10, loIfName, true, true},
		{255, 10, loIfName, true, false},
		{255, 10, loIfName, false, true},
		{255, 10, loIfName, false, false},
		{255, 10, ethIfName, true, true},
		{255, 10, ethIfName, true, false},
		{255, 10, ethIfName, false, true},
		{255, 10, ethIfName, false, false},
		{255, 30, loIfName, true, true},
		{255, 30, loIfName, true, false},
		{255, 30, loIfName, false, true},
		{255, 30, loIfName, false, false},
		{255, 30, ethIfName, true, true},
		{255, 30, ethIfName, true, false},
		{255, 30, ethIfName, false, true},
		{255, 30, ethIfName, false, false},
		{550, 10, loIfName, true, true},
		{550, 10, loIfName, true, false},
		{550, 10, loIfName, false, true},
		{550, 10, loIfName, false, false},
		{550, 10, ethIfName, true, true},
		{550, 10, ethIfName, true, false},
		{550, 10, ethIfName, false, true},
		{550, 10, ethIfName, false, false},
		{550, 30, loIfName, true, true},
		{550, 30, loIfName, true, false},
		{550, 30, loIfName, false, true},
		{550, 30, loIfName, false, false},
		{550, 30, ethIfName, true, true},
		{550, 30, ethIfName, true, false},
		{550, 30, ethIfName, false, true},
		{550, 30, ethIfName, false, false},
	}

	for _, c := range cases {
		g, err := util.GetDefaultGenerator()
		if err != nil {
			return err
		}

		g.SetLinuxCgroupsPath(cgroups.AbsCgroupPath)
		g.SetLinuxResourcesNetworkClassID(c.classid)
		g.AddLinuxResourcesNetworkPriorities(c.ifName, c.prio)

		if !c.withNetNs {
			g.RemoveLinuxNamespace("network")
		}
		if !c.withUserNs {
			g.RemoveLinuxNamespace("user")
		}

		err = util.RuntimeOutsideValidate(g, t, util.ValidateLinuxResourcesNetwork)
		if err != nil {
			return err
		}
	}

	return nil
}

func main() {
	if "linux" != runtime.GOOS {
		util.Fatal(fmt.Errorf("linux-specific cgroup test"))
	}

	if err := testNetworkCgroups(); err != nil {
		util.Fatal(err)
	}
}
<|endoftext|>


<commit_before>package golapa

import (
	"fmt"
	"bytes"
	"text/template"
	"os"
	"net/smtp"
	"crypto/tls"
	_ "crypto/sha512"
	"net"
	"time"
	"log"
	"encoding/base64"
)

type Message struct {
	Sender string
	To	[]string
	Subject string
	Body	string
}

func (msg *Message) ToMessage() ([]byte) {
	var b bytes.Buffer
	t := time.Now()
	curDateTime := t.UTC().Format(time.RFC1123Z)

	b.Write([]byte(fmt.Sprintf("Sender: %s\r\n", msg.Sender)))
	b.Write([]byte(fmt.Sprintf("Reply-To: %s\r\n", msg.Sender)))
	b.Write([]byte(fmt.Sprintf("Subject: =?utf-8?B?%s?=\r\n", base64.URLEncoding.EncodeToString([]byte(msg.Subject)))))
	b.Write([]byte(fmt.Sprintf("Date: %s\r\n", curDateTime)))
	b.Write([]byte(fmt.Sprintf("From: %s\r\n", msg.Sender)))
	b.Write([]byte(fmt.Sprintf("To: %s\r\n", msg.To)))

	return b.Bytes()
}

type Emailer interface {
	Send(msg *Message)
}

type StandardEmailer struct {
	TemplatePath	string
}

func (se *StandardEmailer) Send(msg *Message) {
	host := os.Getenv("EMAIL_HOST")
	password := os.Getenv("EMAIL_PASSWORD")
	from := os.Getenv("EMAIL_FROM")
	ssl := os.Getenv("EMAIL_SSL")

	hostPart, _, err := net.SplitHostPort(host)
	if err != nil {
		log.Print(err.Error())
	}
	auth := smtp.PlainAuth("", from, password, hostPart)

	if ssl == "" {
		err = smtp.SendMail(host, auth, from, msg.To, msg.ToMessage())
		if err != nil {
			log.Print(err.Error())
		}
	}

	if ssl == "true" {
		err = se.SendMailSSL(host, auth, from, msg.To, msg.ToMessage())
		if err != nil {
			log.Print(err.Error())
		}
	}
}

func (se *StandardEmailer) CreateMessage(sender string, subject string, to string, name string, email string) (*Message, error) {
	tpl, err := template.ParseFiles(fmt.Sprintf("%s/email.tpl", se.TemplatePath))
	if err != nil {
		return nil, err
	}
	kv := map[string]string{}
	kv["name"] = name
	kv["email"] = email

	var body bytes.Buffer
	tpl.Execute(&body, kv)
	log.Print(string(body.Bytes()))
	return &Message{
		Sender: sender,
		To: to,
		Subject: subject,
		Body: string(body.Bytes()),
	}, nil
}

func (se *StandardEmailer) SendMailSSL(host string, auth smtp.Auth, from string, recipients []string, msg []byte) (error) {
	conn, err := tls.Dial("tcp", host, nil)
	if err != nil {
		log.Print(err.Error())
		return err
	}

	c, err := smtp.NewClient(conn, host)
	if err != nil {
		log.Print(err.Error())
		return err
	}

	err = c.Auth(auth)
	if err != nil {
		log.Print(err.Error())
		return err
	}

	err = c.Mail(from)
	if err != nil {
		log.Print(err.Error())
		return err
	}

	for _, addr := range recipients {
		err = c.Rcpt(addr)
		if err != nil {
			log.Print(err.Error())
			return err
		}
	}

	d, err := c.Data()
	if err != nil {
		log.Print(err.Error())
		return err
	}

	_, err = d.Write(msg)
	if err != nil {
		log.Print(err.Error())
		return err
	}

	err = d.Close()
	if err != nil {
		log.Print(err.Error())
		return err
	}

	c.Quit()
	return nil
}
<commit_msg>String slice bugfix<commit_after>package golapa

import (
	"fmt"
	"bytes"
	"text/template"
	"os"
	"net/smtp"
	"crypto/tls"
	_ "crypto/sha512"
	"net"
	"time"
	"log"
	"encoding/base64"
)

type Message struct {
	Sender string
	To	string
	Subject string
	Body	string
}

func (msg *Message) ToMessage() ([]byte) {
	var b bytes.Buffer
	t := time.Now()
	curDateTime := t.UTC().Format(time.RFC1123Z)

	b.Write([]byte(fmt.Sprintf("Sender: %s\r\n", msg.Sender)))
	b.Write([]byte(fmt.Sprintf("Reply-To: %s\r\n", msg.Sender)))
	b.Write([]byte(fmt.Sprintf("Subject: =?utf-8?B?%s?=\r\n", base64.URLEncoding.EncodeToString([]byte(msg.Subject)))))
	b.Write([]byte(fmt.Sprintf("Date: %s\r\n", curDateTime)))
	b.Write([]byte(fmt.Sprintf("From: %s\r\n", msg.Sender)))
	b.Write([]byte(fmt.Sprintf("To: %s\r\n", msg.To)))

	return b.Bytes()
}

type Emailer interface {
	Send(msg *Message)
}

type StandardEmailer struct {
	TemplatePath	string
}

func (se *StandardEmailer) Send(msg *Message) {
	host := os.Getenv("EMAIL_HOST")
	password := os.Getenv("EMAIL_PASSWORD")
	from := os.Getenv("EMAIL_FROM")
	ssl := os.Getenv("EMAIL_SSL")

	hostPart, _, err := net.SplitHostPort(host)
	if err != nil {
		log.Print(err.Error())
	}
	auth := smtp.PlainAuth("", from, password, hostPart)

	if ssl == "" {
		err = smtp.SendMail(host, auth, from, msg.To, msg.ToMessage())
		if err != nil {
			log.Print(err.Error())
		}
	}

	if ssl == "true" {
		err = se.SendMailSSL(host, auth, from, msg.To, msg.ToMessage())
		if err != nil {
			log.Print(err.Error())
		}
	}
}

func (se *StandardEmailer) CreateMessage(sender string, subject string, to string, name string, email string) (*Message, error) {
	tpl, err := template.ParseFiles(fmt.Sprintf("%s/email.tpl", se.TemplatePath))
	if err != nil {
		return nil, err
	}
	kv := map[string]string{}
	kv["name"] = name
	kv["email"] = email

	var body bytes.Buffer
	tpl.Execute(&body, kv)
	log.Print(string(body.Bytes()))
	return &Message{
		Sender: sender,
		To: to,
		Subject: subject,
		Body: string(body.Bytes()),
	}, nil
}

func (se *StandardEmailer) SendMailSSL(host string, auth smtp.Auth, from string, recipients []string, msg []byte) (error) {
	conn, err := tls.Dial("tcp", host, nil)
	if err != nil {
		log.Print(err.Error())
		return err
	}

	c, err := smtp.NewClient(conn, host)
	if err != nil {
		log.Print(err.Error())
		return err
	}

	err = c.Auth(auth)
	if err != nil {
		log.Print(err.Error())
		return err
	}

	err = c.Mail(from)
	if err != nil {
		log.Print(err.Error())
		return err
	}

	for _, addr := range recipients {
		err = c.Rcpt(addr)
		if err != nil {
			log.Print(err.Error())
			return err
		}
	}

	d, err := c.Data()
	if err != nil {
		log.Print(err.Error())
		return err
	}

	_, err = d.Write(msg)
	if err != nil {
		log.Print(err.Error())
		return err
	}

	err = d.Close()
	if err != nil {
		log.Print(err.Error())
		return err
	}

	c.Quit()
	return nil
}
<|endoftext|>


<commit_before>// Copyright (c) 2020 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package tworandomchoices

import (
	"math/rand"
	"time"

	"go.uber.org/yarpc/api/peer"
	"go.uber.org/yarpc/api/transport"
	"go.uber.org/yarpc/peer/abstractlist"
)

type twoRandomChoicesList struct {
	subscribers []*subscriber
	random      *rand.Rand
}

// Option configures the peer list implementation constructor.
type Option interface {
	apply(*options)
}

type options struct{}

// NewImplementation creates a new fewest pending heap
// abstractlist.Implementation.
//
// Use this constructor instead of NewList, when wanting to do custom peer
// connection management.
func NewImplementation(opts ...Option) abstractlist.Implementation {
	return newTwoRandomChoicesList(10, rand.NewSource(time.Now().UnixNano()))
}

func newTwoRandomChoicesList(cap int, source rand.Source) *twoRandomChoicesList {
	return &twoRandomChoicesList{
		subscribers: make([]*subscriber, 0, cap),
		random:      rand.New(source),
	}
}

func (l *twoRandomChoicesList) Add(peer peer.StatusPeer, _ peer.Identifier) abstractlist.Subscriber {
	index := len(l.subscribers)
	l.subscribers = append(l.subscribers, &subscriber{
		index: index,
		peer:  peer,
	})
	return l.subscribers[index]
}

func (l *twoRandomChoicesList) Remove(peer peer.StatusPeer, _ peer.Identifier, ps abstractlist.Subscriber) {
	sub, ok := ps.(*subscriber)
	if !ok || len(l.subscribers) == 0 {
		return
	}
	index := sub.index
	last := len(l.subscribers) - 1
	l.subscribers[index] = l.subscribers[last]
	l.subscribers[index].index = index
	l.subscribers = l.subscribers[0:last]
}

func (l *twoRandomChoicesList) Choose(_ *transport.Request) peer.StatusPeer {
	numSubs := len(l.subscribers)
	if numSubs == 0 {
		return nil
	}
	if numSubs == 1 {
		return l.subscribers[0].peer
	}
	i := l.random.Intn(numSubs)
	j := i + 1 + l.random.Intn(numSubs-1)
	if j >= numSubs {
		j -= numSubs
	}
	if l.pending(i) > l.pending(j) {
		i = j
	}
	return l.subscribers[i].peer
}

func (l *twoRandomChoicesList) pending(index int) int {
	return l.subscribers[index].pending
}

type subscriber struct {
	index   int
	peer    peer.StatusPeer
	pending int
}

var _ abstractlist.Subscriber = (*subscriber)(nil)

func (s *subscriber) UpdatePendingRequestCount(pendingRequestCount int) {
	s.pending = pendingRequestCount
}
<commit_msg>Fixed: tworandomchoices now handle its own locking (#2008)<commit_after>// Copyright (c) 2020 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package tworandomchoices

import (
	"math/rand"
	"sync"
	"time"

	"go.uber.org/yarpc/api/peer"
	"go.uber.org/yarpc/api/transport"
	"go.uber.org/yarpc/peer/abstractlist"
)

type twoRandomChoicesList struct {
	subscribers []*subscriber
	random      *rand.Rand

	m sync.RWMutex
}

// Option configures the peer list implementation constructor.
type Option interface {
	apply(*options)
}

type options struct{}

// NewImplementation creates a new fewest pending heap
// abstractlist.Implementation.
//
// Use this constructor instead of NewList, when wanting to do custom peer
// connection management.
func NewImplementation(opts ...Option) abstractlist.Implementation {
	return newTwoRandomChoicesList(10, rand.NewSource(time.Now().UnixNano()))
}

func newTwoRandomChoicesList(cap int, source rand.Source) *twoRandomChoicesList {
	return &twoRandomChoicesList{
		subscribers: make([]*subscriber, 0, cap),
		random:      rand.New(source),
	}
}

func (l *twoRandomChoicesList) Add(peer peer.StatusPeer, _ peer.Identifier) abstractlist.Subscriber {
	l.m.Lock()
	defer l.m.Unlock()

	index := len(l.subscribers)
	l.subscribers = append(l.subscribers, &subscriber{
		index: index,
		peer:  peer,
	})
	return l.subscribers[index]
}

func (l *twoRandomChoicesList) Remove(peer peer.StatusPeer, _ peer.Identifier, ps abstractlist.Subscriber) {
	l.m.Lock()
	defer l.m.Unlock()

	sub, ok := ps.(*subscriber)
	if !ok || len(l.subscribers) == 0 {
		return
	}
	index := sub.index
	last := len(l.subscribers) - 1
	l.subscribers[index] = l.subscribers[last]
	l.subscribers[index].index = index
	l.subscribers = l.subscribers[0:last]
}

func (l *twoRandomChoicesList) Choose(_ *transport.Request) peer.StatusPeer {
	l.m.RLock()
	defer l.m.RUnlock()

	numSubs := len(l.subscribers)
	if numSubs == 0 {
		return nil
	}
	if numSubs == 1 {
		return l.subscribers[0].peer
	}
	i := l.random.Intn(numSubs)
	j := i + 1 + l.random.Intn(numSubs-1)
	if j >= numSubs {
		j -= numSubs
	}
	if l.pending(i) > l.pending(j) {
		i = j
	}
	return l.subscribers[i].peer
}

func (l *twoRandomChoicesList) pending(index int) int {
	return l.subscribers[index].pending
}

type subscriber struct {
	index   int
	peer    peer.StatusPeer
	pending int
}

var _ abstractlist.Subscriber = (*subscriber)(nil)

func (s *subscriber) UpdatePendingRequestCount(pendingRequestCount int) {
	s.pending = pendingRequestCount
}
<|endoftext|>