<commit_before>package pgmgr

import (
	"fmt"
	"io/ioutil"
	"os/exec"
	"path"
	"path/filepath"
	"strings"
	"testing"
	"time"
)

const (
	testDBName      = "pgmgr_testdb"
	migrationFolder = "/tmp/migrations/"
	dumpFile        = "/tmp/pgmgr_dump.sql"
)

func globalConfig() *Config {
	return &Config{
		Database:        testDBName,
		Host:            "localhost",
		Port:            5432,
		DumpFile:        dumpFile,
		MigrationFolder: migrationFolder,
		MigrationTable:  "schema_migrations",
	}
}

func TestCreate(t *testing.T) {
	dropDB(t)

	if err := Create(globalConfig()); err != nil {
		t.Log(err)
		t.Fatal("Could not create database")
	}

	// if we can't remove that db, it couldn't have been created by us above.
	if err := dropDB(t); err != nil {
		t.Fatal("database doesn't seem to have been created!")
	}
}

func TestDrop(t *testing.T) {
	if err := createDB(t); err != nil {
		t.Fatal("createdb failed: ", err)
	}

	if err := Drop(globalConfig()); err != nil {
		t.Log(err)
		t.Fatal("Could not drop database")
	}

	if err := createDB(t); err != nil {
		t.Fatal("database doesn't seem to have been dropped!")
	}
}

func TestDump(t *testing.T) {
	resetDB(t)
	psqlMustExec(t, `CREATE TABLE bars (bar_id INTEGER);`)
	psqlMustExec(t, `INSERT INTO bars (bar_id) VALUES (123), (456);`)
	psqlMustExec(t, `CREATE TABLE foos (foo_id INTEGER);`)
	psqlMustExec(t, `INSERT INTO foos (foo_id) VALUES (789);`)

	c := globalConfig()
	err := Dump(c)

	if err != nil {
		t.Log(err)
		t.Fatal("Could not dump database to file")
	}

	file, err := ioutil.ReadFile(dumpFile)
	if err != nil {
		t.Log(err)
		t.Fatal("Could not read dump")
	}

	if !strings.Contains(string(file), "CREATE TABLE bars") {
		t.Fatal("dump does not contain the table definition")
	}

	if !strings.Contains(string(file), "123") {
		t.Fatal("dump does not contain the table data when --seed-tables is not specified")
	}

	c.SeedTables = append(c.SeedTables, "foos")
	err = Dump(c)

	if err != nil {
		t.Log(err)
		t.Fatal("Could not dump database to file")
	}

	file, err = ioutil.ReadFile(dumpFile)
	if err != nil {
		t.Log(err)
		t.Fatal("Could not read dump")
	}

	if strings.Contains(string(file), "123") {
		t.Fatal("dump contains table data for non-seed tables, when --seed-tables was given")
	}

	if !strings.Contains(string(file), "789") {
		t.Fatal("dump does not contain table data for seed tables, when --seed-tables was given")
	}
}

func TestLoad(t *testing.T) {
	resetDB(t)

	ioutil.WriteFile(dumpFile, []byte(`
		CREATE TABLE foos (foo_id INTEGER);
		INSERT INTO foos (foo_id) VALUES (1), (2), (3);
	`), 0644)

	err := Load(globalConfig())

	if err != nil {
		t.Log(err)
		t.Fatal("Could not load database from file")
	}

	psqlMustExec(t, `SELECT * FROM foos;`)
}

func TestInitialize(t *testing.T) {
	config := globalConfig()

	// Default config should create public.schema_migrations
	resetDB(t)

	if err := Initialize(config); err != nil {
		t.Fatal("Initialize failed:", err)
	}

	if err := Initialize(config); err != nil {
		t.Fatal("Initialize was not safe to run twice:", err)
	}

	psqlMustExec(t, `SELECT * FROM public.schema_migrations;`)

	// If we specify a table, it should create public.<table_name>
	resetDB(t)
	config.MigrationTable = "applied_migrations"

	if err := Initialize(config); err != nil {
		t.Fatal("Initialize failed: ", err)
	}

	psqlMustExec(t, `SELECT * FROM public.applied_migrations;`)

	// If we specify a schema-qualified table, the schema should be
	// created if it does not yet exist.
	resetDB(t)
	config.MigrationTable = "pgmgr.applied_migrations"
	if err := Initialize(config); err != nil {
		t.Fatal("Initialize failed: ", err)
	}

	psqlMustExec(t, `SELECT * FROM pgmgr.applied_migrations`)

	// If we specify a schema-qualified table, and the schema already existed,
	// that's fine too.
	resetDB(t)
	psqlMustExec(t, `CREATE SCHEMA pgmgr;`)
	if err := Initialize(config); err != nil {
		t.Fatal("Initialize failed: ", err)
	}

	psqlMustExec(t, `SELECT * FROM pgmgr.applied_migrations`)
}

func TestVersion(t *testing.T) {
	resetDB(t)

	version, err := Version(globalConfig())

	if err != nil {
		t.Log(err)
		t.Fatal("Could not fetch version info")
	}

	if version != -1 {
		t.Fatal("expected version to be -1 before table exists, got", version)
	}

	Initialize(globalConfig())
	psqlMustExec(t, `INSERT INTO schema_migrations (version) VALUES (1);`)

	version, err = Version(globalConfig())
	if version != 1 {
		t.Fatal("expected version to be 1, got", version)
	}
}

func TestColumnTypeString(t *testing.T) {
	resetDB(t)

	config := globalConfig()
	config.ColumnType = "string"
	Initialize(config)

	psqlMustExec(t, `INSERT INTO schema_migrations (version) VALUES ('20150910120933');`)
	version, err := Version(config)
	if err != nil {
		t.Fatal(err)
	}

	if version != 20150910120933 {
		t.Fatal("expected version to be 20150910120933, got", version)
	}
}

func TestMigrate(t *testing.T) {
	resetDB(t)
	clearMigrationFolder(t)

	// add our first migration
	writeMigration(t, "002_this_is_a_migration.up.sql", `
		CREATE TABLE foos (foo_id INTEGER);
		INSERT INTO foos (foo_id) VALUES (1), (2), (3);
	`)

	writeMigration(t, "002_this_is_a_migration.down.sql", `DROP TABLE foos;`)

	err := Migrate(globalConfig())

	if err != nil {
		t.Log(err)
		t.Fatal("Migrations failed to run.")
	}

	// test simple idempotency
	err = Migrate(globalConfig())
	if err != nil {
		t.Log(err)
		t.Fatal("Running migrations again was not idempotent!")
	}

	psqlMustExec(t, `SELECT * FROM foos;`)

	// add a new migration with an older version, as if another dev's branch was merged in
	writeMigration(t, "001_this_is_an_older_migration.up.sql", `
		CREATE TABLE bars (bar_id INTEGER);
		INSERT INTO bars (bar_id) VALUES (4), (5), (6);
	`)

	err = Migrate(globalConfig())
	if err != nil {
		t.Log(err)
		t.Fatal("Could not apply second migration!")
	}

	psqlMustExec(t, `SELECT * FROM bars;`)

	// Make a filename that would match a vim .swp file
	writeMigration(t, ".003_this_is_an_older_migration.up.sql.swp", `
		CREATE TABLE baz (baz_id INTEGER);
		INSERT INTO baz (baz_id) VALUES (4), (5), (6);
	`)

	err = Migrate(globalConfig())
	if err != nil {
		t.Log(err)
		t.Fatal("Could not apply third migration!")
	}

	psqlMustNotExec(t, `SELECT * FROM baz;`)

	// rollback the initial migration, since it has the latest version
	err = Rollback(globalConfig())

	if err := psqlExec(t, `SELECT * FROM foos;`); err == nil {
		t.Fatal("Should not have been able to select from foos table")
	}

	v, err := Version(globalConfig())
	if err != nil || v != 1 {
		t.Log(err)
		t.Fatal("Rollback did not reset version! Still on version ", v)
	}
}

func TestMigrateColumnTypeString(t *testing.T) {
	resetDB(t)
	clearMigrationFolder(t)

	config := globalConfig()
	config.ColumnType = "string"

	// migrate up
	writeMigration(t, "20150910120933_some_migration.up.sql", `
		CREATE TABLE foos (foo_id INTEGER);
		INSERT INTO foos (foo_id) VALUES (1), (2), (3);
	`)

	err := Migrate(config)
	if err != nil {
		t.Fatal(err)
	}

	v, err := Version(config)
	if err != nil {
		t.Fatal(err)
	}

	if v != 20150910120933 {
		t.Fatal("Expected version 20150910120933 after migration, got", v)
	}

	// migrate down
	writeMigration(t, "20150910120933_some_migration.down.sql", `DROP TABLE foos;`)

	err = Rollback(config)
	if err != nil {
		t.Fatal(err)
	}

	v, err = Version(config)
	if err != nil {
		t.Fatal(err)
	}

	if v != -1 {
		t.Fatal("Expected version -1 after rollback, got", v)
	}
}

func TestMigrateNoTransaction(t *testing.T) {
	resetDB(t)
	clearMigrationFolder(t)

	// CREATE INDEX CONCURRENTLY can not run inside a transaction, so we can assert
	// that no transaction was used by verifying it ran successfully.
	writeMigration(t, "001_create_foos.up.sql", `CREATE TABLE foos (foo_id INTEGER);`)
	writeMigration(t, "002_index_foos.no_txn.up.sql", `CREATE INDEX CONCURRENTLY idx_foo_id ON foos(foo_id);`)

	err := Migrate(globalConfig())
	if err != nil {
		t.Fatal(err)
	}
}

func TestMigrateCustomMigrationTable(t *testing.T) {
	resetDB(t)
	clearMigrationFolder(t)
	writeMigration(t, "001_create_foos.up.sql", `CREATE TABLE foos (foo_id INTEGER);`)

	config := globalConfig()
	config.MigrationTable = "pgmgr.migrations"
	if err := Migrate(config); err != nil {
		t.Fatal(err)
	}

	v, err := Version(config)
	if err != nil {
		t.Fatal(err)
	}

	if v != 1 {
		t.Fatal("Expected version 1, got ", v)
	}
}

func TestCreateMigration(t *testing.T) {
	clearMigrationFolder(t)

	assertFileExists := func(filename string) {
		err := testSh(t, "stat", []string{filepath.Join(migrationFolder, filename)})
		if err != nil {
			t.Fatal(err)
		}
	}

	expectedVersion := time.Now().Unix()
	err := CreateMigration(globalConfig(), "new_migration", false)
	if err != nil {
		t.Fatal(err)
	}

	assertFileExists(fmt.Sprint(expectedVersion, "_new_migration.up.sql"))
	assertFileExists(fmt.Sprint(expectedVersion, "_new_migration.down.sql"))

	expectedStringVersion := time.Now().Format(datetimeFormat)
	config := globalConfig()
	config.Format = "datetime"
	err = CreateMigration(config, "rails_style", false)
	if err != nil {
		t.Fatal(err)
	}

	assertFileExists(fmt.Sprint(expectedStringVersion, "_rails_style.up.sql"))
	assertFileExists(fmt.Sprint(expectedStringVersion, "_rails_style.down.sql"))

	err = CreateMigration(config, "create_index", true)
	if err != nil {
		t.Fatal(err)
	}

	assertFileExists(fmt.Sprint(expectedStringVersion, "_create_index.no_txn.up.sql"))
	assertFileExists(fmt.Sprint(expectedStringVersion, "_create_index.no_txn.down.sql"))
}

// redundant, but I'm also lazy
func testSh(t *testing.T, command string, args []string) error {
	c := exec.Command(command, args...)
	output, err := c.CombinedOutput()
	t.Log(string(output))
	if err != nil {
		return err
	}

	return nil
}

func psqlExec(t *testing.T, statement string) error {
	return testSh(t, "psql", []string{"-d", testDBName, "-c", statement})
}

func psqlMustExec(t *testing.T, statement string) {
	err := psqlExec(t, statement)
	if err != nil {
		t.Fatalf("Failed to execute statement: '%s': %s", statement, err)
	}
}

func psqlMustNotExec(t *testing.T, statement string) {
	err := psqlExec(t, statement)

	// If there is no error, the statement successfully executed.
	// We don't want that to happen.
	if err == nil {
		t.Fatalf("SQL statement executed when it should not have: '%s'", statement)
	}
}

func resetDB(t *testing.T) {
	if err := dropDB(t); err != nil {
		t.Fatal("dropdb failed: ", err)
	}

	if err := createDB(t); err != nil {
		t.Fatal("createdb failed: ", err)
	}
}

func dropDB(t *testing.T) error {
	return testSh(t, "dropdb", []string{testDBName})
}

func createDB(t *testing.T) error {
	return testSh(t, "createdb", []string{testDBName})
}

func clearMigrationFolder(t *testing.T) {
	testSh(t, "rm", []string{"-r", migrationFolder})
	testSh(t, "mkdir", []string{migrationFolder})
}

func writeMigration(t *testing.T, name, contents string) {
	filename := path.Join(migrationFolder, name)
	err := ioutil.WriteFile(filename, []byte(contents), 0644)
	if err != nil {
		t.Fatalf("Failed to write %s: %s", filename, err)
	}
}
<commit_msg>Change error message<commit_after>package pgmgr

import (
	"fmt"
	"io/ioutil"
	"os/exec"
	"path"
	"path/filepath"
	"strings"
	"testing"
	"time"
)

const (
	testDBName      = "pgmgr_testdb"
	migrationFolder = "/tmp/migrations/"
	dumpFile        = "/tmp/pgmgr_dump.sql"
)

func globalConfig() *Config {
	return &Config{
		Database:        testDBName,
		Host:            "localhost",
		Port:            5432,
		DumpFile:        dumpFile,
		MigrationFolder: migrationFolder,
		MigrationTable:  "schema_migrations",
	}
}

func TestCreate(t *testing.T) {
	dropDB(t)

	if err := Create(globalConfig()); err != nil {
		t.Log(err)
		t.Fatal("Could not create database")
	}

	// if we can't remove that db, it couldn't have been created by us above.
	if err := dropDB(t); err != nil {
		t.Fatal("database doesn't seem to have been created!")
	}
}

func TestDrop(t *testing.T) {
	if err := createDB(t); err != nil {
		t.Fatal("createdb failed: ", err)
	}

	if err := Drop(globalConfig()); err != nil {
		t.Log(err)
		t.Fatal("Could not drop database")
	}

	if err := createDB(t); err != nil {
		t.Fatal("database doesn't seem to have been dropped!")
	}
}

func TestDump(t *testing.T) {
	resetDB(t)
	psqlMustExec(t, `CREATE TABLE bars (bar_id INTEGER);`)
	psqlMustExec(t, `INSERT INTO bars (bar_id) VALUES (123), (456);`)
	psqlMustExec(t, `CREATE TABLE foos (foo_id INTEGER);`)
	psqlMustExec(t, `INSERT INTO foos (foo_id) VALUES (789);`)

	c := globalConfig()
	err := Dump(c)

	if err != nil {
		t.Log(err)
		t.Fatal("Could not dump database to file")
	}

	file, err := ioutil.ReadFile(dumpFile)
	if err != nil {
		t.Log(err)
		t.Fatal("Could not read dump")
	}

	if !strings.Contains(string(file), "CREATE TABLE bars") {
		t.Fatal("dump does not contain the table definition")
	}

	if !strings.Contains(string(file), "123") {
		t.Fatal("dump does not contain the table data when --seed-tables is not specified")
	}

	c.SeedTables = append(c.SeedTables, "foos")
	err = Dump(c)

	if err != nil {
		t.Log(err)
		t.Fatal("Could not dump database to file")
	}

	file, err = ioutil.ReadFile(dumpFile)
	if err != nil {
		t.Log(err)
		t.Fatal("Could not read dump")
	}

	if strings.Contains(string(file), "123") {
		t.Fatal("dump contains table data for non-seed tables, when --seed-tables was given")
	}

	if !strings.Contains(string(file), "789") {
		t.Fatal("dump does not contain table data for seed tables, when --seed-tables was given")
	}
}

func TestLoad(t *testing.T) {
	resetDB(t)

	ioutil.WriteFile(dumpFile, []byte(`
		CREATE TABLE foos (foo_id INTEGER);
		INSERT INTO foos (foo_id) VALUES (1), (2), (3);
	`), 0644)

	err := Load(globalConfig())

	if err != nil {
		t.Log(err)
		t.Fatal("Could not load database from file")
	}

	psqlMustExec(t, `SELECT * FROM foos;`)
}

func TestInitialize(t *testing.T) {
	config := globalConfig()

	// Default config should create public.schema_migrations
	resetDB(t)

	if err := Initialize(config); err != nil {
		t.Fatal("Initialize failed:", err)
	}

	if err := Initialize(config); err != nil {
		t.Fatal("Initialize was not safe to run twice:", err)
	}

	psqlMustExec(t, `SELECT * FROM public.schema_migrations;`)

	// If we specify a table, it should create public.<table_name>
	resetDB(t)
	config.MigrationTable = "applied_migrations"

	if err := Initialize(config); err != nil {
		t.Fatal("Initialize failed: ", err)
	}

	psqlMustExec(t, `SELECT * FROM public.applied_migrations;`)

	// If we specify a schema-qualified table, the schema should be
	// created if it does not yet exist.
	resetDB(t)
	config.MigrationTable = "pgmgr.applied_migrations"
	if err := Initialize(config); err != nil {
		t.Fatal("Initialize failed: ", err)
	}

	psqlMustExec(t, `SELECT * FROM pgmgr.applied_migrations`)

	// If we specify a schema-qualified table, and the schema already existed,
	// that's fine too.
	resetDB(t)
	psqlMustExec(t, `CREATE SCHEMA pgmgr;`)
	if err := Initialize(config); err != nil {
		t.Fatal("Initialize failed: ", err)
	}

	psqlMustExec(t, `SELECT * FROM pgmgr.applied_migrations`)
}

func TestVersion(t *testing.T) {
	resetDB(t)

	version, err := Version(globalConfig())

	if err != nil {
		t.Log(err)
		t.Fatal("Could not fetch version info")
	}

	if version != -1 {
		t.Fatal("expected version to be -1 before table exists, got", version)
	}

	Initialize(globalConfig())
	psqlMustExec(t, `INSERT INTO schema_migrations (version) VALUES (1);`)

	version, err = Version(globalConfig())
	if version != 1 {
		t.Fatal("expected version to be 1, got", version)
	}
}

func TestColumnTypeString(t *testing.T) {
	resetDB(t)

	config := globalConfig()
	config.ColumnType = "string"
	Initialize(config)

	psqlMustExec(t, `INSERT INTO schema_migrations (version) VALUES ('20150910120933');`)
	version, err := Version(config)
	if err != nil {
		t.Fatal(err)
	}

	if version != 20150910120933 {
		t.Fatal("expected version to be 20150910120933, got", version)
	}
}

func TestMigrate(t *testing.T) {
	resetDB(t)
	clearMigrationFolder(t)

	// add our first migration
	writeMigration(t, "002_this_is_a_migration.up.sql", `
		CREATE TABLE foos (foo_id INTEGER);
		INSERT INTO foos (foo_id) VALUES (1), (2), (3);
	`)

	writeMigration(t, "002_this_is_a_migration.down.sql", `DROP TABLE foos;`)

	err := Migrate(globalConfig())

	if err != nil {
		t.Log(err)
		t.Fatal("Migrations failed to run.")
	}

	// test simple idempotency
	err = Migrate(globalConfig())
	if err != nil {
		t.Log(err)
		t.Fatal("Running migrations again was not idempotent!")
	}

	psqlMustExec(t, `SELECT * FROM foos;`)

	// add a new migration with an older version, as if another dev's branch was merged in
	writeMigration(t, "001_this_is_an_older_migration.up.sql", `
		CREATE TABLE bars (bar_id INTEGER);
		INSERT INTO bars (bar_id) VALUES (4), (5), (6);
	`)

	err = Migrate(globalConfig())
	if err != nil {
		t.Log(err)
		t.Fatal("Could not apply second migration!")
	}

	psqlMustExec(t, `SELECT * FROM bars;`)

	// Make a filename that would match a vim .swp file
	writeMigration(t, ".003_this_is_an_older_migration.up.sql.swp", `
		CREATE TABLE baz (baz_id INTEGER);
		INSERT INTO baz (baz_id) VALUES (4), (5), (6);
	`)

	err = Migrate(globalConfig())
	if err != nil {
		t.Log(err)
		t.Fatal("Migration returned an error instead of being skipped!")
	}

	psqlMustNotExec(t, `SELECT * FROM baz;`)

	// rollback the initial migration, since it has the latest version
	err = Rollback(globalConfig())

	if err := psqlExec(t, `SELECT * FROM foos;`); err == nil {
		t.Fatal("Should not have been able to select from foos table")
	}

	v, err := Version(globalConfig())
	if err != nil || v != 1 {
		t.Log(err)
		t.Fatal("Rollback did not reset version! Still on version ", v)
	}
}

func TestMigrateColumnTypeString(t *testing.T) {
	resetDB(t)
	clearMigrationFolder(t)

	config := globalConfig()
	config.ColumnType = "string"

	// migrate up
	writeMigration(t, "20150910120933_some_migration.up.sql", `
		CREATE TABLE foos (foo_id INTEGER);
		INSERT INTO foos (foo_id) VALUES (1), (2), (3);
	`)

	err := Migrate(config)
	if err != nil {
		t.Fatal(err)
	}

	v, err := Version(config)
	if err != nil {
		t.Fatal(err)
	}

	if v != 20150910120933 {
		t.Fatal("Expected version 20150910120933 after migration, got", v)
	}

	// migrate down
	writeMigration(t, "20150910120933_some_migration.down.sql", `DROP TABLE foos;`)

	err = Rollback(config)
	if err != nil {
		t.Fatal(err)
	}

	v, err = Version(config)
	if err != nil {
		t.Fatal(err)
	}

	if v != -1 {
		t.Fatal("Expected version -1 after rollback, got", v)
	}
}

func TestMigrateNoTransaction(t *testing.T) {
	resetDB(t)
	clearMigrationFolder(t)

	// CREATE INDEX CONCURRENTLY can not run inside a transaction, so we can assert
	// that no transaction was used by verifying it ran successfully.
	writeMigration(t, "001_create_foos.up.sql", `CREATE TABLE foos (foo_id INTEGER);`)
	writeMigration(t, "002_index_foos.no_txn.up.sql", `CREATE INDEX CONCURRENTLY idx_foo_id ON foos(foo_id);`)

	err := Migrate(globalConfig())
	if err != nil {
		t.Fatal(err)
	}
}

func TestMigrateCustomMigrationTable(t *testing.T) {
	resetDB(t)
	clearMigrationFolder(t)
	writeMigration(t, "001_create_foos.up.sql", `CREATE TABLE foos (foo_id INTEGER);`)

	config := globalConfig()
	config.MigrationTable = "pgmgr.migrations"
	if err := Migrate(config); err != nil {
		t.Fatal(err)
	}

	v, err := Version(config)
	if err != nil {
		t.Fatal(err)
	}

	if v != 1 {
		t.Fatal("Expected version 1, got ", v)
	}
}

func TestCreateMigration(t *testing.T) {
	clearMigrationFolder(t)

	assertFileExists := func(filename string) {
		err := testSh(t, "stat", []string{filepath.Join(migrationFolder, filename)})
		if err != nil {
			t.Fatal(err)
		}
	}

	expectedVersion := time.Now().Unix()
	err := CreateMigration(globalConfig(), "new_migration", false)
	if err != nil {
		t.Fatal(err)
	}

	assertFileExists(fmt.Sprint(expectedVersion, "_new_migration.up.sql"))
	assertFileExists(fmt.Sprint(expectedVersion, "_new_migration.down.sql"))

	expectedStringVersion := time.Now().Format(datetimeFormat)
	config := globalConfig()
	config.Format = "datetime"
	err = CreateMigration(config, "rails_style", false)
	if err != nil {
		t.Fatal(err)
	}

	assertFileExists(fmt.Sprint(expectedStringVersion, "_rails_style.up.sql"))
	assertFileExists(fmt.Sprint(expectedStringVersion, "_rails_style.down.sql"))

	err = CreateMigration(config, "create_index", true)
	if err != nil {
		t.Fatal(err)
	}

	assertFileExists(fmt.Sprint(expectedStringVersion, "_create_index.no_txn.up.sql"))
	assertFileExists(fmt.Sprint(expectedStringVersion, "_create_index.no_txn.down.sql"))
}

// redundant, but I'm also lazy
func testSh(t *testing.T, command string, args []string) error {
	c := exec.Command(command, args...)
	output, err := c.CombinedOutput()
	t.Log(string(output))
	if err != nil {
		return err
	}

	return nil
}

func psqlExec(t *testing.T, statement string) error {
	return testSh(t, "psql", []string{"-d", testDBName, "-c", statement})
}

func psqlMustExec(t *testing.T, statement string) {
	err := psqlExec(t, statement)
	if err != nil {
		t.Fatalf("Failed to execute statement: '%s': %s", statement, err)
	}
}

func psqlMustNotExec(t *testing.T, statement string) {
	err := psqlExec(t, statement)

	// If there is no error, the statement successfully executed.
	// We don't want that to happen.
	if err == nil {
		t.Fatalf("SQL statement executed when it should not have: '%s'", statement)
	}
}

func resetDB(t *testing.T) {
	if err := dropDB(t); err != nil {
		t.Fatal("dropdb failed: ", err)
	}

	if err := createDB(t); err != nil {
		t.Fatal("createdb failed: ", err)
	}
}

func dropDB(t *testing.T) error {
	return testSh(t, "dropdb", []string{testDBName})
}

func createDB(t *testing.T) error {
	return testSh(t, "createdb", []string{testDBName})
}

func clearMigrationFolder(t *testing.T) {
	testSh(t, "rm", []string{"-r", migrationFolder})
	testSh(t, "mkdir", []string{migrationFolder})
}

func writeMigration(t *testing.T, name, contents string) {
	filename := path.Join(migrationFolder, name)
	err := ioutil.WriteFile(filename, []byte(contents), 0644)
	if err != nil {
		t.Fatalf("Failed to write %s: %s", filename, err)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"database/sql"
	"fmt"
	"github.com/lib/pq"
	"log"
	"os"
	"os/signal"
	"runtime/debug"
	"syscall"
	"time"
)

func WorkerExtractPgerror(err error) (*string, error) {
	pgerr, ok := err.(pq.PGError)
	if ok {
		msg := pgerr.Get('M')
		return &msg, nil
	}
	if err.Error() == "driver: bad connection" {
		msg := "could not connect to database"
		return &msg, nil
	}
	return nil, err
}

// WorkerCoerceType returns a coerced version of the raw
// database value in, which we get from scanning into
// interface{}s. We expect queries from the following
// Postgres types to result in the following return values:
// [Postgres] -> [Go: in] -> [Go: WorkerCoerceType'd]
// text          []byte      string
// ???
func WorkerCoerceType(in interface{}) interface{} {
	switch in := in.(type) {
	case []byte:
		return string(in)
	default:
		return in
	}
}

// WorkerQuery queries the pin db at pinDbUrl and updates the
// passed pin according to the results/errors. System errors
// are returned.
func WorkerQuery(p *Pin, pinDbUrl string) error {
	log.Printf("worker.query.start pin_id=%s", p.Id)
	applicationName := fmt.Sprintf("pgpin.pin.%s", p.Id)
	pinDbConn := fmt.Sprintf("%s?application_name=%s&statement_timeout=%d&connect_timeout=%d",
		pinDbUrl, applicationName, DataPinStatementTimeout/time.Millisecond, DataConnectTimeout/time.Millisecond)
	pinDb, err := sql.Open("postgres", pinDbConn)
	if err != nil {
		p.ResultsError, err = WorkerExtractPgerror(err)
		return err
	}
	resultsRows, err := pinDb.Query(p.Query)
	if err != nil {
		p.ResultsError, err = WorkerExtractPgerror(err)
		return err
	}
	defer func() { Must(resultsRows.Close()) }()
	resultsFieldsData, err := resultsRows.Columns()
	if err != nil {
		p.ResultsError, err = WorkerExtractPgerror(err)
		return err
	}
	resultsRowsData := make([][]interface{}, 0)
	resultsRowsSeen := 0
	for resultsRows.Next() {
		resultsRowsSeen += 1
		if resultsRowsSeen > DataPinResultsRowsMax {
			message := "too many rows in query results"
			p.ResultsError = &message
			return nil
		}
		resultsRowData := make([]interface{}, len(resultsFieldsData))
		resultsRowPointers := make([]interface{}, len(resultsFieldsData))
		for i, _ := range resultsRowData {
			resultsRowPointers[i] = &resultsRowData[i]
		}
		err := resultsRows.Scan(resultsRowPointers...)
		if err != nil {
			p.ResultsError, err = WorkerExtractPgerror(err)
			return err
		}
		for i, _ := range resultsRowData {
			resultsRowData[i] = WorkerCoerceType(resultsRowData[i])
		}
		resultsRowsData = append(resultsRowsData, resultsRowData)
	}
	err = resultsRows.Err()
	if err != nil {
		p.ResultsError, err = WorkerExtractPgerror(err)
		return err
	}
	p.ResultsFields = MustNewPgJson(resultsFieldsData)
	p.ResultsRows = MustNewPgJson(resultsRowsData)
	log.Printf("worker.query.finish pin_id=%s", p.Id)
	return nil
}

// WorkerProcess performs a processes an update on the given
// pin, running its query against its db and updating the
// system database accordingly. User-caused errors are
// reflected in the updated pin record and will not cause a
// returned error. System-caused errors are returned.
func WorkerProcess(p *Pin) error {
	log.Printf("worker.process.start pin_id=%s", p.Id)
	pinDbUrl, err := DataPinDbUrl(p)
	if err != nil {
		return err
	}
	startedAt := time.Now()
	p.QueryStartedAt = &startedAt
	err = WorkerQuery(p, pinDbUrl)
	if err != nil {
		return err
	}
	finishedAt := time.Now()
	p.QueryFinishedAt = &finishedAt
	err = DataPinUpdate(p)
	if err != nil {
		return err
	}
	log.Printf("worker.process.finish pin_id=%s", p.Id)
	return nil
}

// WorkerTick processes 1 pending pin, if such a pin is
// available. It returns true iff a pin is successfully processed.
func WorkerTick() (bool, error) {
	log.Printf("worker.tick.start")
	p, err := DataPinReserve()
	if err != nil {
		return false, err
	}
	if p != nil {
		log.Printf("worker.tick.found pin_id=%s", p.Id)
		err = WorkerProcess(p)
		if err != nil {
			return false, err
		}
		err = DataPinRelease(p)
		if err != nil {
			return false, err
		}
		return true, nil
	}
	return false, nil
}

// WorkerTrap returns a chanel that will be populated when
// an INT or TERM signals is received.
func WorkerTrap() chan bool {
	sig := make(chan os.Signal, 1)
	done := make(chan bool, 1)
	signal.Notify(sig, syscall.SIGINT, syscall.SIGTERM)
	go func() {
		<-sig
		log.Printf("worker.trap")
		done <- true
	}()
	return done
}

var WorkerCooloff = time.Millisecond * 500

func WorkerCheckPanic() {
	err := recover()
	if err != nil {
		log.Printf("worker.panic: %s", err)
		log.Print(string(debug.Stack()))
		time.Sleep(WorkerCooloff)
	}
}

func WorkerHandleError(err error) {
	log.Printf("worker.error %s", err.Error())
	time.Sleep(WorkerCooloff)
}

func WorkerCheckExit(done chan bool) {
	select {
	case <-done:
		log.Printf("worker.exit")
		os.Exit(0)
	default:
	}
}

func WorkerLoop(done chan bool) {
	defer WorkerCheckPanic()
	processed, err := WorkerTick()
	if err != nil {
		WorkerHandleError(err)
	}
	WorkerCheckExit(done)
	if err == nil && !processed {
		time.Sleep(WorkerCooloff)
	}
}

func WorkerStart() {
	log.Printf("worker.start")
	DataStart()
	done := WorkerTrap()
	for {
		WorkerLoop(done)
	}
}
<commit_msg>Spelling<commit_after>package main

import (
	"database/sql"
	"fmt"
	"github.com/lib/pq"
	"log"
	"os"
	"os/signal"
	"runtime/debug"
	"syscall"
	"time"
)

func WorkerExtractPgerror(err error) (*string, error) {
	pgerr, ok := err.(pq.PGError)
	if ok {
		msg := pgerr.Get('M')
		return &msg, nil
	}
	if err.Error() == "driver: bad connection" {
		msg := "could not connect to database"
		return &msg, nil
	}
	return nil, err
}

// WorkerCoerceType returns a coerced version of the raw
// database value in, which we get from scanning into
// interface{}s. We expect queries from the following
// Postgres types to result in the following return values:
// [Postgres] -> [Go: in] -> [Go: WorkerCoerceType'd]
// text          []byte      string
// ???
func WorkerCoerceType(in interface{}) interface{} {
	switch in := in.(type) {
	case []byte:
		return string(in)
	default:
		return in
	}
}

// WorkerQuery queries the pin db at pinDbUrl and updates the
// passed pin according to the results/errors. System errors
// are returned.
func WorkerQuery(p *Pin, pinDbUrl string) error {
	log.Printf("worker.query.start pin_id=%s", p.Id)
	applicationName := fmt.Sprintf("pgpin.pin.%s", p.Id)
	pinDbConn := fmt.Sprintf("%s?application_name=%s&statement_timeout=%d&connect_timeout=%d",
		pinDbUrl, applicationName, DataPinStatementTimeout/time.Millisecond, DataConnectTimeout/time.Millisecond)
	pinDb, err := sql.Open("postgres", pinDbConn)
	if err != nil {
		p.ResultsError, err = WorkerExtractPgerror(err)
		return err
	}
	resultsRows, err := pinDb.Query(p.Query)
	if err != nil {
		p.ResultsError, err = WorkerExtractPgerror(err)
		return err
	}
	defer func() { Must(resultsRows.Close()) }()
	resultsFieldsData, err := resultsRows.Columns()
	if err != nil {
		p.ResultsError, err = WorkerExtractPgerror(err)
		return err
	}
	resultsRowsData := make([][]interface{}, 0)
	resultsRowsSeen := 0
	for resultsRows.Next() {
		resultsRowsSeen += 1
		if resultsRowsSeen > DataPinResultsRowsMax {
			message := "too many rows in query results"
			p.ResultsError = &message
			return nil
		}
		resultsRowData := make([]interface{}, len(resultsFieldsData))
		resultsRowPointers := make([]interface{}, len(resultsFieldsData))
		for i, _ := range resultsRowData {
			resultsRowPointers[i] = &resultsRowData[i]
		}
		err := resultsRows.Scan(resultsRowPointers...)
		if err != nil {
			p.ResultsError, err = WorkerExtractPgerror(err)
			return err
		}
		for i, _ := range resultsRowData {
			resultsRowData[i] = WorkerCoerceType(resultsRowData[i])
		}
		resultsRowsData = append(resultsRowsData, resultsRowData)
	}
	err = resultsRows.Err()
	if err != nil {
		p.ResultsError, err = WorkerExtractPgerror(err)
		return err
	}
	p.ResultsFields = MustNewPgJson(resultsFieldsData)
	p.ResultsRows = MustNewPgJson(resultsRowsData)
	log.Printf("worker.query.finish pin_id=%s", p.Id)
	return nil
}

// WorkerProcess performs a processes an update on the given
// pin, running its query against its db and updating the
// system database accordingly. User-caused errors are
// reflected in the updated pin record and will not cause a
// returned error. System-caused errors are returned.
func WorkerProcess(p *Pin) error {
	log.Printf("worker.process.start pin_id=%s", p.Id)
	pinDbUrl, err := DataPinDbUrl(p)
	if err != nil {
		return err
	}
	startedAt := time.Now()
	p.QueryStartedAt = &startedAt
	err = WorkerQuery(p, pinDbUrl)
	if err != nil {
		return err
	}
	finishedAt := time.Now()
	p.QueryFinishedAt = &finishedAt
	err = DataPinUpdate(p)
	if err != nil {
		return err
	}
	log.Printf("worker.process.finish pin_id=%s", p.Id)
	return nil
}

// WorkerTick processes 1 pending pin, if such a pin is
// available. It returns true iff a pin is successfully processed.
func WorkerTick() (bool, error) {
	log.Printf("worker.tick.start")
	p, err := DataPinReserve()
	if err != nil {
		return false, err
	}
	if p != nil {
		log.Printf("worker.tick.found pin_id=%s", p.Id)
		err = WorkerProcess(p)
		if err != nil {
			return false, err
		}
		err = DataPinRelease(p)
		if err != nil {
			return false, err
		}
		return true, nil
	}
	return false, nil
}

// WorkerTrap returns a channel that will be populated when
// an INT or TERM signals is received.
func WorkerTrap() chan bool {
	sig := make(chan os.Signal, 1)
	done := make(chan bool, 1)
	signal.Notify(sig, syscall.SIGINT, syscall.SIGTERM)
	go func() {
		<-sig
		log.Printf("worker.trap")
		done <- true
	}()
	return done
}

var WorkerCooloff = time.Millisecond * 500

func WorkerCheckPanic() {
	err := recover()
	if err != nil {
		log.Printf("worker.panic: %s", err)
		log.Print(string(debug.Stack()))
		time.Sleep(WorkerCooloff)
	}
}

func WorkerHandleError(err error) {
	log.Printf("worker.error %s", err.Error())
	time.Sleep(WorkerCooloff)
}

func WorkerCheckExit(done chan bool) {
	select {
	case <-done:
		log.Printf("worker.exit")
		os.Exit(0)
	default:
	}
}

func WorkerLoop(done chan bool) {
	defer WorkerCheckPanic()
	processed, err := WorkerTick()
	if err != nil {
		WorkerHandleError(err)
	}
	WorkerCheckExit(done)
	if err == nil && !processed {
		time.Sleep(WorkerCooloff)
	}
}

func WorkerStart() {
	log.Printf("worker.start")
	DataStart()
	done := WorkerTrap()
	for {
		WorkerLoop(done)
	}
}
<|endoftext|>


<commit_before>package dgroup

import (
	"context"
	"os"
	"os/signal"
	"sort"
	"syscall"

	"github.com/pkg/errors"

	"github.com/datawire/ambassador/pkg/derrgroup"
	"github.com/datawire/ambassador/pkg/dlog"
)

// Group is a wrapper around golang.org/x/sync/errgroup.Group (err, a
// fork of errgroup) that:
//  - handles SIGINT and SIGTERM
type Group struct {
	hardCtx       context.Context
	softCtx       context.Context
	loggerFactory func(name string) dlog.Logger
	inner         *derrgroup.Group
}

func logGoroutines(printf func(format string, args ...interface{}), list map[string]derrgroup.GoroutineState) {
	printf("  goroutine shutdown status:")
	names := make([]string, 0, len(list))
	nameWidth := 0
	for name := range list {
		names = append(names, name)
		if len(name) > nameWidth {
			nameWidth = len(name)
		}
	}
	sort.Strings(names)
	for _, name := range names {
		printf("    %-*s: %s", nameWidth, name, list[name])
	}
}

// NewGroup returns a new Group.
func NewGroup(ctx context.Context, loggerFactory func(name string) dlog.Logger) *Group {
	sigs := make(chan os.Signal, 1)
	signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM)

	hardCtx, hardCancel := context.WithCancel(ctx)
	softCtx, softCancel := context.WithCancel(hardCtx)

	ret := &Group{
		hardCtx:       hardCtx,
		softCtx:       softCtx,
		loggerFactory: loggerFactory,
		inner:         derrgroup.NewGroup(softCancel),
	}

	ret.Go("supervisor", func(hardCtx, _ context.Context) error {
		<-softCtx.Done()
		dlog.Infoln(hardCtx, "shutting down...")
		return nil
	})

	ret.Go("signal_handler", func(hardCtx, _ context.Context) error {
		defer func() {
			// If we receive another signal after
			// graceful-shutdown, we should trigger a
			// not-so-graceful shutdown.
			go func() {
				sig := <-sigs
				dlog.Errorln(hardCtx, errors.Errorf("received signal %v", sig))
				errorf := func(fmt string, args ...interface{}) {
					dlog.Errorf(hardCtx, fmt, args...)
				}
				logGoroutines(errorf, ret.List())
				hardCancel()
				// keep logging signals
				for sig := range sigs {
					dlog.Errorln(hardCtx, errors.Errorf("received signal %v", sig))
					logGoroutines(errorf, ret.List())
				}
			}()
		}()

		select {
		case sig := <-sigs:
			return errors.Errorf("received signal %v", sig)
		case <-softCtx.Done():
			return nil
		}
	})

	return ret
}

// Go wraps derrgroup.Group.Go().
//
//  - `softCtx` being canceled should trigger a graceful shutdown
//  - `hardCtx` being canceled should trigger a not-so-graceful shutdown
func (g *Group) Go(name string, fn func(hardCtx, softCtx context.Context) error) {
	g.inner.Go(name, func() error {
		logger := g.loggerFactory(name)
		hardCtx := dlog.WithLogger(g.hardCtx, logger)
		softCtx := dlog.WithLogger(g.softCtx, logger)
		err := fn(hardCtx, softCtx)
		if err == nil {
			logger.Debugln("goroutine exited without error")
		} else {
			logger.Errorln("goroutine exited with error:", err)
		}
		return err
	})
}

// Wait wraps derrgroup.Group.Wait().
func (g *Group) Wait() error {
	ret := g.inner.Wait()
	if ret != nil {
		logGoroutines(g.loggerFactory("shutdown_status").Infof, g.List())
	}
	return ret
}

// List wraps derrgroup.Group.List().
func (g *Group) List() map[string]derrgroup.GoroutineState {
	return g.inner.List()
}
<commit_msg>(from AES) dgroup: Update comment<commit_after>package dgroup

import (
	"context"
	"os"
	"os/signal"
	"sort"
	"syscall"

	"github.com/pkg/errors"

	"github.com/datawire/ambassador/pkg/derrgroup"
	"github.com/datawire/ambassador/pkg/dlog"
)

// Group is a wrapper around
// github.com/datawire/ambassador/pkg/derrgroup.Group that:
//  - handles SIGINT and SIGTERM
//  - manages Context for you
//  - adds hard/soft cancelation
//  - does some minimal logging
type Group struct {
	hardCtx       context.Context
	softCtx       context.Context
	loggerFactory func(name string) dlog.Logger
	inner         *derrgroup.Group
}

func logGoroutines(printf func(format string, args ...interface{}), list map[string]derrgroup.GoroutineState) {
	printf("  goroutine shutdown status:")
	names := make([]string, 0, len(list))
	nameWidth := 0
	for name := range list {
		names = append(names, name)
		if len(name) > nameWidth {
			nameWidth = len(name)
		}
	}
	sort.Strings(names)
	for _, name := range names {
		printf("    %-*s: %s", nameWidth, name, list[name])
	}
}

// NewGroup returns a new Group.
func NewGroup(ctx context.Context, loggerFactory func(name string) dlog.Logger) *Group {
	sigs := make(chan os.Signal, 1)
	signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM)

	hardCtx, hardCancel := context.WithCancel(ctx)
	softCtx, softCancel := context.WithCancel(hardCtx)

	ret := &Group{
		hardCtx:       hardCtx,
		softCtx:       softCtx,
		loggerFactory: loggerFactory,
		inner:         derrgroup.NewGroup(softCancel),
	}

	ret.Go("supervisor", func(hardCtx, _ context.Context) error {
		<-softCtx.Done()
		dlog.Infoln(hardCtx, "shutting down...")
		return nil
	})

	ret.Go("signal_handler", func(hardCtx, _ context.Context) error {
		defer func() {
			// If we receive another signal after
			// graceful-shutdown, we should trigger a
			// not-so-graceful shutdown.
			go func() {
				sig := <-sigs
				dlog.Errorln(hardCtx, errors.Errorf("received signal %v", sig))
				errorf := func(fmt string, args ...interface{}) {
					dlog.Errorf(hardCtx, fmt, args...)
				}
				logGoroutines(errorf, ret.List())
				hardCancel()
				// keep logging signals
				for sig := range sigs {
					dlog.Errorln(hardCtx, errors.Errorf("received signal %v", sig))
					logGoroutines(errorf, ret.List())
				}
			}()
		}()

		select {
		case sig := <-sigs:
			return errors.Errorf("received signal %v", sig)
		case <-softCtx.Done():
			return nil
		}
	})

	return ret
}

// Go wraps derrgroup.Group.Go().
//
//  - `softCtx` being canceled should trigger a graceful shutdown
//  - `hardCtx` being canceled should trigger a not-so-graceful shutdown
func (g *Group) Go(name string, fn func(hardCtx, softCtx context.Context) error) {
	g.inner.Go(name, func() error {
		logger := g.loggerFactory(name)
		hardCtx := dlog.WithLogger(g.hardCtx, logger)
		softCtx := dlog.WithLogger(g.softCtx, logger)
		err := fn(hardCtx, softCtx)
		if err == nil {
			logger.Debugln("goroutine exited without error")
		} else {
			logger.Errorln("goroutine exited with error:", err)
		}
		return err
	})
}

// Wait wraps derrgroup.Group.Wait().
func (g *Group) Wait() error {
	ret := g.inner.Wait()
	if ret != nil {
		logGoroutines(g.loggerFactory("shutdown_status").Infof, g.List())
	}
	return ret
}

// List wraps derrgroup.Group.List().
func (g *Group) List() map[string]derrgroup.GoroutineState {
	return g.inner.List()
}
<|endoftext|>


<commit_before>// Copyright 2015 CoreOS, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package omaha

import (
	"bytes"
	"encoding/xml"
	"fmt"
	"io/ioutil"
	"net/http"
	"sync"
	"testing"

	"github.com/coreos/mantle/Godeps/_workspace/src/github.com/kylelemons/godebug/pretty"
)

func TestServerRequestResponse(t *testing.T) {
	var wg sync.WaitGroup
	defer wg.Wait()

	// make an omaha server
	s, err := NewServer(":0")
	if err != nil {
		t.Fatalf("failed to create omaha server: %v", err)
	}

	wg.Add(1)
	go func() {
		defer wg.Done()
		if err := s.Serve(); err != nil {
			t.Errorf("Serve failed: %v", err)
		}
	}()

	defer s.Stop()

	// make an omaha request
	request := NewRequest()

	// ensures the struct is the same as what appears out of the Decoder in
	// Server's handler
	request.XMLName.Local = "request"

	request.OS.Platform = "CoreOS"

	buf := new(bytes.Buffer)
	enc := xml.NewEncoder(buf)
	enc.Indent("", "\t")
	err = enc.Encode(request)
	if err != nil {
		t.Errorf("failed to marshal request: %v", err)
		return
	}

	// check that server gets the same thing we sent
	rch := s.RequestChan()
	wg.Add(1)
	go func() {
		defer wg.Done()
		sreq, ok := <-rch
		if !ok {
			t.Errorf("failed to get notification from server")
			return
		}

		if diff := pretty.Compare(request, sreq); diff != "" {
			t.Errorf("client request differs from what server got: %v", diff)
		}
	}()

	// send omaha request
	endpoint := fmt.Sprintf("http://%s/v1/update/", s.Addr())
	t.Logf("sending request to %q:\n%s\n", endpoint, buf)
	res, err := http.Post(endpoint, "text/xml", buf)
	if err != nil {
		t.Errorf("failed to post: %v", err)
		return
	}

	defer res.Body.Close()

	if res.StatusCode != 200 {
		t.Errorf("failed to post: %v", res.Status)
		return
	}

	body, _ := ioutil.ReadAll(res.Body)

	t.Logf("got response:\n%s\n", body)

}
<commit_msg>omaha: fix server test to send valid request, validate response<commit_after>// Copyright 2015 CoreOS, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package omaha

import (
	"bytes"
	"encoding/xml"
	"fmt"
	"net/http"
	"sync"
	"testing"
)

func TestServerRequestResponse(t *testing.T) {
	var wg sync.WaitGroup
	defer wg.Wait()

	// make an omaha server
	s, err := NewServer(":0")
	if err != nil {
		t.Fatalf("failed to create omaha server: %v", err)
	}

	wg.Add(1)
	go func() {
		defer wg.Done()
		if err := s.Serve(); err != nil {
			t.Errorf("Serve failed: %v", err)
		}
	}()

	defer s.Stop()

	buf := new(bytes.Buffer)
	enc := xml.NewEncoder(buf)
	enc.Indent("", "\t")
	err = enc.Encode(nilRequest)
	if err != nil {
		t.Errorf("failed to marshal request: %v", err)
		return
	}

	// check that server gets the same thing we sent
	rch := s.RequestChan()
	wg.Add(1)
	go func() {
		defer wg.Done()
		sreq, ok := <-rch
		if !ok {
			t.Errorf("failed to get notification from server")
			return
		}

		if err := compareXML(nilRequest, sreq); err != nil {
			t.Error(err)
		}
	}()

	// send omaha request
	endpoint := fmt.Sprintf("http://%s/v1/update/", s.Addr())
	res, err := http.Post(endpoint, "text/xml", buf)
	if err != nil {
		t.Errorf("failed to post: %v", err)
		return
	}

	defer res.Body.Close()

	if res.StatusCode != 200 {
		t.Errorf("failed to post: %v", res.Status)
		return
	}

	dec := xml.NewDecoder(res.Body)
	sresp := &Response{}
	if err := dec.Decode(sresp); err != nil {
		t.Fatalf("failed to parse body: %v", err)
	}
	if err := compareXML(nilResponse, sresp); err != nil {
		t.Error(err)
	}
}
<|endoftext|>


<commit_before>package negroni

import (
	"bufio"
	"fmt"
	"net"
	"net/http"
)

// ResponseWriter is a wrapper around http.ResponseWriter that provides extra information about
// the response. It is recommended that middleware handlers use this construct to wrap a responsewriter
// if the functionality calls for it.
type ResponseWriter interface {
	http.ResponseWriter
	http.Flusher
	// Status returns the status code of the response or 200 if the response has
	// not been written (as this is the default response code in net/http)
	Status() int
	// Written returns whether or not the ResponseWriter has been written.
	Written() bool
	// Size returns the size of the response body.
	Size() int
	// Before allows for a function to be called before the ResponseWriter has been written to. This is
	// useful for setting headers or any other operations that must happen before a response has been written.
	Before(func(ResponseWriter))
}

type beforeFunc func(ResponseWriter)

// NewResponseWriter creates a ResponseWriter that wraps an http.ResponseWriter
func NewResponseWriter(rw http.ResponseWriter) ResponseWriter {
	nrw := &responseWriter{
		ResponseWriter: rw,
	}

	if _, ok := rw.(http.CloseNotifier); ok {
		return &responseWriterCloseNotifer{nrw}
	}

	return nrw
}

type responseWriter struct {
	http.ResponseWriter
	status      int
	size        int
	beforeFuncs []beforeFunc
}

func (rw *responseWriter) WriteHeader(s int) {
	rw.status = s
	rw.callBefore()
	rw.ResponseWriter.WriteHeader(s)
}

func (rw *responseWriter) Write(b []byte) (int, error) {
	if !rw.Written() {
		// The status will be StatusOK if WriteHeader has not been called yet
		rw.WriteHeader(http.StatusOK)
	}
	size, err := rw.ResponseWriter.Write(b)
	rw.size += size
	return size, err
}

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

func (rw *responseWriter) Size() int {
	return rw.size
}

func (rw *responseWriter) Written() bool {
	return rw.status != 0
}

func (rw *responseWriter) Before(before func(ResponseWriter)) {
	rw.beforeFuncs = append(rw.beforeFuncs, before)
}

func (rw *responseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
	hijacker, ok := rw.ResponseWriter.(http.Hijacker)
	if !ok {
		return nil, nil, fmt.Errorf("the ResponseWriter doesn't support the Hijacker interface")
	}
	return hijacker.Hijack()
}

func (rw *responseWriter) callBefore() {
	for i := len(rw.beforeFuncs) - 1; i >= 0; i-- {
		rw.beforeFuncs[i](rw)
	}
}

func (rw *responseWriter) Flush() {
	flusher, ok := rw.ResponseWriter.(http.Flusher)
	if ok {
		if !rw.Written() {
			// The status will be StatusOK if WriteHeader has not been called yet
			rw.WriteHeader(http.StatusOK)
		}
		flusher.Flush()
	}
}

type responseWriterCloseNotifer struct {
	*responseWriter
}

func (rw *responseWriterCloseNotifer) CloseNotify() <-chan bool {
	return rw.ResponseWriter.(http.CloseNotifier).CloseNotify()
}
<commit_msg>Update documentation of ResponseWriter.Status() to match implementation<commit_after>package negroni

import (
	"bufio"
	"fmt"
	"net"
	"net/http"
)

// ResponseWriter is a wrapper around http.ResponseWriter that provides extra information about
// the response. It is recommended that middleware handlers use this construct to wrap a responsewriter
// if the functionality calls for it.
type ResponseWriter interface {
	http.ResponseWriter
	http.Flusher
	// Status returns the status code of the response or 0 if the response has
	// not been written
	Status() int
	// Written returns whether or not the ResponseWriter has been written.
	Written() bool
	// Size returns the size of the response body.
	Size() int
	// Before allows for a function to be called before the ResponseWriter has been written to. This is
	// useful for setting headers or any other operations that must happen before a response has been written.
	Before(func(ResponseWriter))
}

type beforeFunc func(ResponseWriter)

// NewResponseWriter creates a ResponseWriter that wraps an http.ResponseWriter
func NewResponseWriter(rw http.ResponseWriter) ResponseWriter {
	nrw := &responseWriter{
		ResponseWriter: rw,
	}

	if _, ok := rw.(http.CloseNotifier); ok {
		return &responseWriterCloseNotifer{nrw}
	}

	return nrw
}

type responseWriter struct {
	http.ResponseWriter
	status      int
	size        int
	beforeFuncs []beforeFunc
}

func (rw *responseWriter) WriteHeader(s int) {
	rw.status = s
	rw.callBefore()
	rw.ResponseWriter.WriteHeader(s)
}

func (rw *responseWriter) Write(b []byte) (int, error) {
	if !rw.Written() {
		// The status will be StatusOK if WriteHeader has not been called yet
		rw.WriteHeader(http.StatusOK)
	}
	size, err := rw.ResponseWriter.Write(b)
	rw.size += size
	return size, err
}

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

func (rw *responseWriter) Size() int {
	return rw.size
}

func (rw *responseWriter) Written() bool {
	return rw.status != 0
}

func (rw *responseWriter) Before(before func(ResponseWriter)) {
	rw.beforeFuncs = append(rw.beforeFuncs, before)
}

func (rw *responseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
	hijacker, ok := rw.ResponseWriter.(http.Hijacker)
	if !ok {
		return nil, nil, fmt.Errorf("the ResponseWriter doesn't support the Hijacker interface")
	}
	return hijacker.Hijack()
}

func (rw *responseWriter) callBefore() {
	for i := len(rw.beforeFuncs) - 1; i >= 0; i-- {
		rw.beforeFuncs[i](rw)
	}
}

func (rw *responseWriter) Flush() {
	flusher, ok := rw.ResponseWriter.(http.Flusher)
	if ok {
		if !rw.Written() {
			// The status will be StatusOK if WriteHeader has not been called yet
			rw.WriteHeader(http.StatusOK)
		}
		flusher.Flush()
	}
}

type responseWriterCloseNotifer struct {
	*responseWriter
}

func (rw *responseWriterCloseNotifer) CloseNotify() <-chan bool {
	return rw.ResponseWriter.(http.CloseNotifier).CloseNotify()
}
<|endoftext|>


<commit_before>package markup

import (
	"bytes"
	"fmt"
	"github.com/yvasiyarov/swagger/parser"
	"log"
	"os"
	"path"
	"sort"
	"strings"
)

type Markup interface {
	sectionHeader(level int, text string) string
	bulletedItem(level int, text string) string
	numberedItem(level int, text string) string
	anchor(anchorName string) string
	link(anchorName, linkText string) string
	tableHeader(tableTitle string) string
	tableHeaderRow(args ...string) string
	tableRow(args ...string) string
	tableFooter() string
}

func GenerateMarkup(parser *parser.Parser, markup Markup, outputSpec *string, defaultFileExtension string) {
	var filename string
	if *outputSpec == "" {
		filename = path.Join("./", "API", defaultFileExtension)
	} else {
		filename = path.Join(*outputSpec)
	}
	fd, err := os.Create(filename)
	if err != nil {
		log.Fatalf("Can not create document file: %v\n", err)
	}
	defer fd.Close()

	var buf bytes.Buffer

	/***************************************************************
	* Overall API
	***************************************************************/
	buf.WriteString(markup.sectionHeader(1, parser.Listing.Infos.Title))
	buf.WriteString(fmt.Sprintf("%s\n\n", parser.Listing.Infos.Description))

	/***************************************************************
	* Table of Contents (List of Sub-APIs)
	***************************************************************/
	buf.WriteString("Table of Contents\n\n")
	subApiKeys, subApiKeyIndex := alphabeticalKeysOfSubApis(parser.Listing.Apis)
	for _, subApiKey := range subApiKeys {
		buf.WriteString(markup.numberedItem(1, markup.link(subApiKey, parser.Listing.Apis[subApiKeyIndex[subApiKey]].Description)))
	}
	buf.WriteString("\n")

	for _, apiKey := range alphabeticalKeysOfApiDeclaration(parser.TopLevelApis) {

		apiDescription := parser.TopLevelApis[apiKey]
		/***************************************************************
		* Sub-API Specifications
		***************************************************************/
		buf.WriteString(markup.anchor(apiKey))
		buf.WriteString(markup.sectionHeader(2, apiKey))

		buf.WriteString(markup.tableHeader(""))
		buf.WriteString(markup.tableHeaderRow("Specification", "Value"))
		buf.WriteString(markup.tableRow("Resource Path", apiDescription.ResourcePath))
		buf.WriteString(markup.tableRow("API Version", apiDescription.ApiVersion))
		buf.WriteString(markup.tableRow("BasePath for the API", apiDescription.BasePath))
		buf.WriteString(markup.tableRow("Consumes", strings.Join(apiDescription.Consumes, ", ")))
		buf.WriteString(markup.tableRow("Produces", strings.Join(apiDescription.Produces, ", ")))
		buf.WriteString(markup.tableFooter())

		/***************************************************************
		* Sub-API Operations (Summary)
		***************************************************************/
		buf.WriteString("\n")
		buf.WriteString(markup.sectionHeader(3, "Operations"))
		buf.WriteString("\n")

		buf.WriteString(markup.tableHeader(""))
		buf.WriteString(markup.tableHeaderRow("Resource Path", "Operation", "Description"))
		for _, subapi := range apiDescription.Apis {
			for _, op := range subapi.Operations {
				pathString := strings.Replace(strings.Replace(subapi.Path, "{", "\\{", -1), "}", "\\}", -1)
				buf.WriteString(markup.tableRow(pathString, markup.link(op.Nickname, op.HttpMethod), op.Summary))
			}
		}
		buf.WriteString(markup.tableFooter())
		buf.WriteString("\n")

		/***************************************************************
		* Sub-API Operations (Details)
		***************************************************************/
		for _, subapi := range apiDescription.Apis {
			for _, op := range subapi.Operations {
				buf.WriteString("\n")
				operationString := fmt.Sprintf("%s (%s)", strings.Replace(strings.Replace(subapi.Path, "{", "\\{", -1), "}", "\\}", -1), op.HttpMethod)
				buf.WriteString(markup.anchor(op.Nickname))
				buf.WriteString(markup.sectionHeader(4, "API: "+operationString))
				buf.WriteString("\n\n" + op.Summary + "\n\n\n")

				if len(op.Parameters) > 0 {
					buf.WriteString(markup.tableHeader(""))
					buf.WriteString(markup.tableHeaderRow("Param Name", "Param Type", "Data Type", "Description", "Required?"))
					for _, param := range op.Parameters {
						isRequired := ""
						if param.Required {
							isRequired = "Yes"
						}
						buf.WriteString(markup.tableRow(param.Name, param.ParamType, modelText(markup, param.DataType), param.Description, isRequired))
					}
					buf.WriteString(markup.tableFooter())
				}

				if len(op.ResponseMessages) > 0 {
					buf.WriteString(markup.tableHeader(""))
					buf.WriteString(markup.tableHeaderRow("Code", "Message", "Model"))
					for _, msg := range op.ResponseMessages {
						buf.WriteString(markup.tableRow(fmt.Sprintf("%v", msg.Code), msg.Message, modelText(markup, msg.ResponseModel)))
					}
					buf.WriteString(markup.tableFooter())
				}
			}
		}
		buf.WriteString("\n")

		/***************************************************************
		* Models
		***************************************************************/
		buf.WriteString("\n")
		buf.WriteString(markup.sectionHeader(3, "Models"))
		buf.WriteString("\n")

		for _, modelKey := range alphabeticalKeysOfModels(apiDescription.Models) {
			model := apiDescription.Models[modelKey]
			buf.WriteString(markup.anchor(modelKey))
			buf.WriteString(markup.sectionHeader(4, shortModelName(modelKey)))
			buf.WriteString(markup.tableHeader(""))
			buf.WriteString(markup.tableHeaderRow("Field Name (alphabetical)", "Field Type", "Description"))
			for _, fieldName := range alphabeticalKeysOfFields(model.Properties) {
				fieldProps := model.Properties[fieldName]
				buf.WriteString(markup.tableRow(fieldName, fieldProps.Type, fieldProps.Description))
			}
			buf.WriteString(markup.tableFooter())
		}
		buf.WriteString("\n")

	}

	fd.WriteString(buf.String())
}

func shortModelName(longModelName string) string {
	parts := strings.Split(longModelName, ".")
	return parts[len(parts)-1]
}

func modelText(markup Markup, fullyQualifiedModelName string) string {
	shortName := shortModelName(fullyQualifiedModelName)
	result := shortName
	if fullyQualifiedModelName != shortName {
		result = markup.link(fullyQualifiedModelName, shortName)
	}
	return result
}

func alphabeticalKeysOfSubApis(refs []*parser.ApiRef) ([]string, map[string]int) {
	index := map[string]int{}
	keys := make([]string, len(refs))
	for i, ref := range refs {
		subApiKey := ref.Path[1:]
		keys[i] = subApiKey
		index[subApiKey] = i
	}
	sort.Strings(keys)
	return keys, index
}
func alphabeticalKeysOfApiDeclaration(m map[string]*parser.ApiDeclaration) []string {
	keys := make([]string, len(m))
	i := 0
	for key, _ := range m {
		keys[i] = key
		i++
	}
	sort.Strings(keys)
	return keys
}
func alphabeticalKeysOfModels(m map[string]*parser.Model) []string {
	keys := make([]string, len(m))
	i := 0
	for key, _ := range m {
		keys[i] = key
		i++
	}
	sort.Strings(keys)
	return keys
}
func alphabeticalKeysOfFields(m map[string]*parser.ModelProperty) []string {
	keys := make([]string, len(m))
	i := 0
	for key, _ := range m {
		keys[i] = key
		i++
	}
	sort.Strings(keys)
	return keys
}
<commit_msg>remove default extension from path.Join call<commit_after>package markup

import (
	"bytes"
	"fmt"
	"github.com/yvasiyarov/swagger/parser"
	"log"
	"os"
	"path"
	"sort"
	"strings"
)

type Markup interface {
	sectionHeader(level int, text string) string
	bulletedItem(level int, text string) string
	numberedItem(level int, text string) string
	anchor(anchorName string) string
	link(anchorName, linkText string) string
	tableHeader(tableTitle string) string
	tableHeaderRow(args ...string) string
	tableRow(args ...string) string
	tableFooter() string
}

func GenerateMarkup(parser *parser.Parser, markup Markup, outputSpec *string, defaultFileExtension string) {
	var filename string
	if *outputSpec == "" {
		filename = path.Join("./", "API") + defaultFileExtension
	} else {
		filename = path.Join(*outputSpec)
	}
	fd, err := os.Create(filename)
	if err != nil {
		log.Fatalf("Can not create document file: %v\n", err)
	}
	defer fd.Close()

	var buf bytes.Buffer

	/***************************************************************
	* Overall API
	***************************************************************/
	buf.WriteString(markup.sectionHeader(1, parser.Listing.Infos.Title))
	buf.WriteString(fmt.Sprintf("%s\n\n", parser.Listing.Infos.Description))

	/***************************************************************
	* Table of Contents (List of Sub-APIs)
	***************************************************************/
	buf.WriteString("Table of Contents\n\n")
	subApiKeys, subApiKeyIndex := alphabeticalKeysOfSubApis(parser.Listing.Apis)
	for _, subApiKey := range subApiKeys {
		buf.WriteString(markup.numberedItem(1, markup.link(subApiKey, parser.Listing.Apis[subApiKeyIndex[subApiKey]].Description)))
	}
	buf.WriteString("\n")

	for _, apiKey := range alphabeticalKeysOfApiDeclaration(parser.TopLevelApis) {

		apiDescription := parser.TopLevelApis[apiKey]
		/***************************************************************
		* Sub-API Specifications
		***************************************************************/
		buf.WriteString(markup.anchor(apiKey))
		buf.WriteString(markup.sectionHeader(2, apiKey))

		buf.WriteString(markup.tableHeader(""))
		buf.WriteString(markup.tableHeaderRow("Specification", "Value"))
		buf.WriteString(markup.tableRow("Resource Path", apiDescription.ResourcePath))
		buf.WriteString(markup.tableRow("API Version", apiDescription.ApiVersion))
		buf.WriteString(markup.tableRow("BasePath for the API", apiDescription.BasePath))
		buf.WriteString(markup.tableRow("Consumes", strings.Join(apiDescription.Consumes, ", ")))
		buf.WriteString(markup.tableRow("Produces", strings.Join(apiDescription.Produces, ", ")))
		buf.WriteString(markup.tableFooter())

		/***************************************************************
		* Sub-API Operations (Summary)
		***************************************************************/
		buf.WriteString("\n")
		buf.WriteString(markup.sectionHeader(3, "Operations"))
		buf.WriteString("\n")

		buf.WriteString(markup.tableHeader(""))
		buf.WriteString(markup.tableHeaderRow("Resource Path", "Operation", "Description"))
		for _, subapi := range apiDescription.Apis {
			for _, op := range subapi.Operations {
				pathString := strings.Replace(strings.Replace(subapi.Path, "{", "\\{", -1), "}", "\\}", -1)
				buf.WriteString(markup.tableRow(pathString, markup.link(op.Nickname, op.HttpMethod), op.Summary))
			}
		}
		buf.WriteString(markup.tableFooter())
		buf.WriteString("\n")

		/***************************************************************
		* Sub-API Operations (Details)
		***************************************************************/
		for _, subapi := range apiDescription.Apis {
			for _, op := range subapi.Operations {
				buf.WriteString("\n")
				operationString := fmt.Sprintf("%s (%s)", strings.Replace(strings.Replace(subapi.Path, "{", "\\{", -1), "}", "\\}", -1), op.HttpMethod)
				buf.WriteString(markup.anchor(op.Nickname))
				buf.WriteString(markup.sectionHeader(4, "API: "+operationString))
				buf.WriteString("\n\n" + op.Summary + "\n\n\n")

				if len(op.Parameters) > 0 {
					buf.WriteString(markup.tableHeader(""))
					buf.WriteString(markup.tableHeaderRow("Param Name", "Param Type", "Data Type", "Description", "Required?"))
					for _, param := range op.Parameters {
						isRequired := ""
						if param.Required {
							isRequired = "Yes"
						}
						buf.WriteString(markup.tableRow(param.Name, param.ParamType, modelText(markup, param.DataType), param.Description, isRequired))
					}
					buf.WriteString(markup.tableFooter())
				}

				if len(op.ResponseMessages) > 0 {
					buf.WriteString(markup.tableHeader(""))
					buf.WriteString(markup.tableHeaderRow("Code", "Message", "Model"))
					for _, msg := range op.ResponseMessages {
						buf.WriteString(markup.tableRow(fmt.Sprintf("%v", msg.Code), msg.Message, modelText(markup, msg.ResponseModel)))
					}
					buf.WriteString(markup.tableFooter())
				}
			}
		}
		buf.WriteString("\n")

		/***************************************************************
		* Models
		***************************************************************/
		buf.WriteString("\n")
		buf.WriteString(markup.sectionHeader(3, "Models"))
		buf.WriteString("\n")

		for _, modelKey := range alphabeticalKeysOfModels(apiDescription.Models) {
			model := apiDescription.Models[modelKey]
			buf.WriteString(markup.anchor(modelKey))
			buf.WriteString(markup.sectionHeader(4, shortModelName(modelKey)))
			buf.WriteString(markup.tableHeader(""))
			buf.WriteString(markup.tableHeaderRow("Field Name (alphabetical)", "Field Type", "Description"))
			for _, fieldName := range alphabeticalKeysOfFields(model.Properties) {
				fieldProps := model.Properties[fieldName]
				buf.WriteString(markup.tableRow(fieldName, fieldProps.Type, fieldProps.Description))
			}
			buf.WriteString(markup.tableFooter())
		}
		buf.WriteString("\n")

	}

	fd.WriteString(buf.String())
}

func shortModelName(longModelName string) string {
	parts := strings.Split(longModelName, ".")
	return parts[len(parts)-1]
}

func modelText(markup Markup, fullyQualifiedModelName string) string {
	shortName := shortModelName(fullyQualifiedModelName)
	result := shortName
	if fullyQualifiedModelName != shortName {
		result = markup.link(fullyQualifiedModelName, shortName)
	}
	return result
}

func alphabeticalKeysOfSubApis(refs []*parser.ApiRef) ([]string, map[string]int) {
	index := map[string]int{}
	keys := make([]string, len(refs))
	for i, ref := range refs {
		subApiKey := ref.Path[1:]
		keys[i] = subApiKey
		index[subApiKey] = i
	}
	sort.Strings(keys)
	return keys, index
}
func alphabeticalKeysOfApiDeclaration(m map[string]*parser.ApiDeclaration) []string {
	keys := make([]string, len(m))
	i := 0
	for key, _ := range m {
		keys[i] = key
		i++
	}
	sort.Strings(keys)
	return keys
}
func alphabeticalKeysOfModels(m map[string]*parser.Model) []string {
	keys := make([]string, len(m))
	i := 0
	for key, _ := range m {
		keys[i] = key
		i++
	}
	sort.Strings(keys)
	return keys
}
func alphabeticalKeysOfFields(m map[string]*parser.ModelProperty) []string {
	keys := make([]string, len(m))
	i := 0
	for key, _ := range m {
		keys[i] = key
		i++
	}
	sort.Strings(keys)
	return keys
}
<|endoftext|>


<commit_before>package merkle

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"strings"

	"github.com/aybabtme/epher/thash"
)

type Store interface {
	PutNode(node Node) error
	GetNode(thash.Sum) (Node, bool, error)

	PutBlob(r io.Reader, done func() thash.Sum) error
	GetBlob(thash.Sum) (io.ReadCloser, error)
	InfoBlob(thash.Sum) (int64, bool, error)
}

type Option func(*config)

type config struct {
	HashType thash.Type `json:"hash_type"`
	BlobSize int64      `json:"blob_size"`
}

func newConfig(opts []Option) *config {
	def := &config{
		HashType: thash.Blake2B512,
		BlobSize: 4 << 20, // 4MiB
	}
	for _, o := range opts {
		o(def)
	}
	return def
}

func WithBlobSize(sz int64) Option      { return func(opts *config) { opts.BlobSize = sz } }
func WithHashType(ht thash.Type) Option { return func(opts *config) { opts.HashType = ht } }

func Build(r io.Reader, store Store, opts ...Option) (*Tree, thash.Sum, error) {

	config := newConfig(opts)

	buf := bytes.NewBuffer(nil)

	rdbuf := bytes.NewBuffer(nil)

	var bis []blobInfo

	reachedEOF := false
	for !reachedEOF {
		buf.Reset()
		rdbuf.Reset()

		n, err := io.CopyN(rdbuf, r, config.BlobSize)
		if err != nil && err != io.EOF {
			return nil, thash.Sum{}, err
		}
		reachedEOF = (err == io.EOF)
		if n == 0 {
			break
		}

		sum, n, err := copyBlob(config.HashType, buf, rdbuf)
		if err != nil {
			return nil, thash.Sum{}, err
		}

		getSum := func() thash.Sum { return sum }

		if err := store.PutBlob(buf, getSum); err != nil {
			return nil, thash.Sum{}, err
		}

		bis = append(bis, blobInfo{sum: sum, size: n})
	}

	tree := newTree(bis)
	if err := tree.persist(store); err != nil {
		return nil, thash.Sum{}, err
	}
	return tree, tree.HashSum, nil
}

func copyBlob(t thash.Type, w io.Writer, r io.Reader) (thash.Sum, int64, error) {

	h := thash.New(t)

	n, err := io.Copy(w, io.TeeReader(r, h))
	if err != nil {
		return thash.Sum{}, n, err
	}

	return thash.MakeSum(h), n, err
}

var errMalformedTree = errors.New("tree is malformed")

// Node is a node in a merkle tree. A node is sufficient
// to retrieve the whole of a merkle tree rooted in this node.
type Node struct {
	Sum        thash.Sum
	Start, End thash.Sum
}

func RetrieveTree(sum thash.Sum, store Store) (*Tree, error) {
	root, found, err := store.GetNode(sum)
	if err != nil {
		return nil, err
	}
	if !found {
		return &Tree{HashSum: sum}, nil
	}
	start, err := RetrieveTree(root.Start, store)
	if err != nil {
		return nil, err
	}
	end, err := RetrieveTree(root.End, store)
	if err != nil {
		return nil, err
	}
	tree := &Tree{Start: start, End: end, HashSum: sum}

	// set the sizes
	err = walk(tree, func(branch *Tree) error {
		branch.SizeByte = branch.Start.SizeByte + branch.End.SizeByte
		return nil
	}, func(leaf *Tree) error {
		size, _, err := store.InfoBlob(leaf.HashSum)
		leaf.SizeByte = size
		return err
	})

	return tree, err
}

// Tree is a concrete merkle tree.
type Tree struct {
	Start *Tree `json:"start"`
	End   *Tree `json:"end"`

	SizeByte int64     `json:"size_byte"`
	HashSum  thash.Sum `json:"hash_sum"`
}

func walk(tree *Tree, onBranch, onLeaf func(*Tree) error) error {
	switch {
	case tree.Start != nil && tree.End != nil:
		if err := walk(tree.Start, onBranch, onLeaf); err != nil {
			return err
		}
		if err := walk(tree.End, onBranch, onLeaf); err != nil {
			return err
		}
		return onBranch(tree)

	case tree.Start == nil && tree.End == nil:
		return onLeaf(tree)

	case tree == nil,
		tree.Start == nil && tree.End != nil, // can't have an end without a start
		tree.Start != nil && tree.End == nil: // we should have been a data node
		return errMalformedTree
	default:
		panic(fmt.Sprintf("unhandled case: %#v", tree))
	}
}

func (tree *Tree) Retrieve(wr io.Writer, store Store) (invalid []*Tree, err error) {
	if wr == nil {
		wr = ioutil.Discard
	}
	return tree.retrieve(wr, store)
}

func (tree *Tree) retrieve(wr io.Writer, store Store) (invalid []*Tree, err error) {
	onBranch := func(branch *Tree) error {
		// verify that this.sum == sum(start.sum, end.sum)
		// then verify:
		//   - start
		//   - end

		// we're a sum of hash sum
		got := sumHashWithTree(branch.Start, branch.End)
		if !branch.HashSum.Equal(got) {
			invalid = append(invalid, branch) // we're invalid
		}
		return nil
	}
	onLeaf := func(leaf *Tree) error {
		rd, err := store.GetBlob(leaf.HashSum)
		if err != nil {
			invalid = []*Tree{leaf}
			return err
		}
		got, _, err := copyBlob(leaf.HashSum.Type, wr, rd)
		if err != nil {
			invalid = []*Tree{leaf}
			return err
		}
		if !leaf.HashSum.Equal(got) {
			invalid = []*Tree{leaf}
			return fmt.Errorf("want sum %v, got %v", leaf.HashSum, got)
		}
		return nil
	}
	return invalid, walk(tree, onBranch, onLeaf)
}

func (tree *Tree) persist(store Store) error {

	onBranch := func(branch *Tree) error {
		if err := branch.Start.persist(store); err != nil {
			return err
		}
		if err := branch.End.persist(store); err != nil {
			return err
		}
		return store.PutNode(Node{
			Sum:   branch.HashSum,
			Start: branch.Start.HashSum,
			End:   branch.End.HashSum,
		})
	}
	onLeaf := func(leaf *Tree) error {
		return nil
	}
	return walk(tree, onBranch, onLeaf)
}

type blobInfo struct {
	sum  thash.Sum
	size int64
}

func newTree(bis []blobInfo) *Tree {
	switch n := len(bis); n {
	case 0: // no data
		return nil
	case 1: // we're a leaf
		bi := bis[0]
		return &Tree{
			HashSum:  bi.sum,
			SizeByte: bi.size,
		}
	default:

		var (
			start = newTree(bis[:n/2])
			end   = newTree(bis[n/2:])
		)
		if start == nil {
			panic("should not be possible")
		}
		if end == nil {
			// we're a tree of odd size
			return start
		}

		// we have two child nodes, so compute the hash sum of their appended hash sums

		return &Tree{
			Start:    start,
			End:      end,
			SizeByte: start.SizeByte + end.SizeByte,
			HashSum:  sumHashWithTree(start, end),
		}
	}
}

func sumHashWithTree(start, end *Tree) thash.Sum {
	appendedSums := start.HashSum.Sum + end.HashSum.Sum
	h := thash.New(start.HashSum.Type)
	_, err := io.Copy(h, strings.NewReader(appendedSums))
	if err != nil {
		panic(err) // should never happen
	}
	return thash.MakeSum(h)
}
<commit_msg>missing close<commit_after>package merkle

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"strings"

	"github.com/aybabtme/epher/thash"
)

type Store interface {
	PutNode(node Node) error
	GetNode(thash.Sum) (Node, bool, error)

	PutBlob(r io.Reader, done func() thash.Sum) error
	GetBlob(thash.Sum) (io.ReadCloser, error)
	InfoBlob(thash.Sum) (int64, bool, error)
}

type Option func(*config)

type config struct {
	HashType thash.Type `json:"hash_type"`
	BlobSize int64      `json:"blob_size"`
}

func newConfig(opts []Option) *config {
	def := &config{
		HashType: thash.Blake2B512,
		BlobSize: 4 << 20, // 4MiB
	}
	for _, o := range opts {
		o(def)
	}
	return def
}

func WithBlobSize(sz int64) Option      { return func(opts *config) { opts.BlobSize = sz } }
func WithHashType(ht thash.Type) Option { return func(opts *config) { opts.HashType = ht } }

func Build(r io.Reader, store Store, opts ...Option) (*Tree, thash.Sum, error) {

	config := newConfig(opts)

	buf := bytes.NewBuffer(nil)

	rdbuf := bytes.NewBuffer(nil)

	var bis []blobInfo

	reachedEOF := false
	for !reachedEOF {
		buf.Reset()
		rdbuf.Reset()

		n, err := io.CopyN(rdbuf, r, config.BlobSize)
		if err != nil && err != io.EOF {
			return nil, thash.Sum{}, err
		}
		reachedEOF = (err == io.EOF)
		if n == 0 {
			break
		}

		sum, n, err := copyBlob(config.HashType, buf, rdbuf)
		if err != nil {
			return nil, thash.Sum{}, err
		}

		getSum := func() thash.Sum { return sum }

		if err := store.PutBlob(buf, getSum); err != nil {
			return nil, thash.Sum{}, err
		}

		bis = append(bis, blobInfo{sum: sum, size: n})
	}

	tree := newTree(bis)
	if err := tree.persist(store); err != nil {
		return nil, thash.Sum{}, err
	}
	return tree, tree.HashSum, nil
}

func copyBlob(t thash.Type, w io.Writer, r io.Reader) (thash.Sum, int64, error) {

	h := thash.New(t)

	n, err := io.Copy(w, io.TeeReader(r, h))
	if err != nil {
		return thash.Sum{}, n, err
	}

	return thash.MakeSum(h), n, err
}

var errMalformedTree = errors.New("tree is malformed")

// Node is a node in a merkle tree. A node is sufficient
// to retrieve the whole of a merkle tree rooted in this node.
type Node struct {
	Sum        thash.Sum
	Start, End thash.Sum
}

func RetrieveTree(sum thash.Sum, store Store) (*Tree, error) {
	root, found, err := store.GetNode(sum)
	if err != nil {
		return nil, err
	}
	if !found {
		return &Tree{HashSum: sum}, nil
	}
	start, err := RetrieveTree(root.Start, store)
	if err != nil {
		return nil, err
	}
	end, err := RetrieveTree(root.End, store)
	if err != nil {
		return nil, err
	}
	tree := &Tree{Start: start, End: end, HashSum: sum}

	// set the sizes
	err = walk(tree, func(branch *Tree) error {
		branch.SizeByte = branch.Start.SizeByte + branch.End.SizeByte
		return nil
	}, func(leaf *Tree) error {
		size, _, err := store.InfoBlob(leaf.HashSum)
		leaf.SizeByte = size
		return err
	})

	return tree, err
}

// Tree is a concrete merkle tree.
type Tree struct {
	Start *Tree `json:"start"`
	End   *Tree `json:"end"`

	SizeByte int64     `json:"size_byte"`
	HashSum  thash.Sum `json:"hash_sum"`
}

func walk(tree *Tree, onBranch, onLeaf func(*Tree) error) error {
	switch {
	case tree.Start != nil && tree.End != nil:
		if err := walk(tree.Start, onBranch, onLeaf); err != nil {
			return err
		}
		if err := walk(tree.End, onBranch, onLeaf); err != nil {
			return err
		}
		return onBranch(tree)

	case tree.Start == nil && tree.End == nil:
		return onLeaf(tree)

	case tree == nil,
		tree.Start == nil && tree.End != nil, // can't have an end without a start
		tree.Start != nil && tree.End == nil: // we should have been a data node
		return errMalformedTree
	default:
		panic(fmt.Sprintf("unhandled case: %#v", tree))
	}
}

func (tree *Tree) Retrieve(wr io.Writer, store Store) (invalid []*Tree, err error) {
	if wr == nil {
		wr = ioutil.Discard
	}
	return tree.retrieve(wr, store)
}

func (tree *Tree) retrieve(wr io.Writer, store Store) (invalid []*Tree, err error) {
	onBranch := func(branch *Tree) error {
		// verify that this.sum == sum(start.sum, end.sum)
		// then verify:
		//   - start
		//   - end

		// we're a sum of hash sum
		got := sumHashWithTree(branch.Start, branch.End)
		if !branch.HashSum.Equal(got) {
			invalid = append(invalid, branch) // we're invalid
		}
		return nil
	}
	onLeaf := func(leaf *Tree) error {
		rd, err := store.GetBlob(leaf.HashSum)
		if err != nil {
			invalid = []*Tree{leaf}
			return err
		}
		defer rd.Close()
		got, _, err := copyBlob(leaf.HashSum.Type, wr, rd)
		if err != nil {
			invalid = []*Tree{leaf}
			return err
		}
		if !leaf.HashSum.Equal(got) {
			invalid = []*Tree{leaf}
			return fmt.Errorf("want sum %v, got %v", leaf.HashSum, got)
		}
		return nil
	}
	return invalid, walk(tree, onBranch, onLeaf)
}

func (tree *Tree) persist(store Store) error {

	onBranch := func(branch *Tree) error {
		if err := branch.Start.persist(store); err != nil {
			return err
		}
		if err := branch.End.persist(store); err != nil {
			return err
		}
		return store.PutNode(Node{
			Sum:   branch.HashSum,
			Start: branch.Start.HashSum,
			End:   branch.End.HashSum,
		})
	}
	onLeaf := func(leaf *Tree) error {
		return nil
	}
	return walk(tree, onBranch, onLeaf)
}

type blobInfo struct {
	sum  thash.Sum
	size int64
}

func newTree(bis []blobInfo) *Tree {
	switch n := len(bis); n {
	case 0: // no data
		return nil
	case 1: // we're a leaf
		bi := bis[0]
		return &Tree{
			HashSum:  bi.sum,
			SizeByte: bi.size,
		}
	default:

		var (
			start = newTree(bis[:n/2])
			end   = newTree(bis[n/2:])
		)
		if start == nil {
			panic("should not be possible")
		}
		if end == nil {
			// we're a tree of odd size
			return start
		}

		// we have two child nodes, so compute the hash sum of their appended hash sums

		return &Tree{
			Start:    start,
			End:      end,
			SizeByte: start.SizeByte + end.SizeByte,
			HashSum:  sumHashWithTree(start, end),
		}
	}
}

func sumHashWithTree(start, end *Tree) thash.Sum {
	appendedSums := start.HashSum.Sum + end.HashSum.Sum
	h := thash.New(start.HashSum.Type)
	_, err := io.Copy(h, strings.NewReader(appendedSums))
	if err != nil {
		panic(err) // should never happen
	}
	return thash.MakeSum(h)
}
<|endoftext|>


<commit_before>package mesh

import (
	"errors"

	"github.com/keybase/saltpack"
	"github.com/keybase/saltpack/basic"
)

// PeerInfo holds the information wire needs to connect to a remote peer
type PeerInfo struct {
	ID        string   `json:"id"`
	Addresses []string `json:"addresses"`
	PublicKey [32]byte `json:"public_key"`
	Signature []byte   `json:"signature"`
}

type peerInfoClean struct {
	ID        string   `json:"id"`
	Addresses []string `json:"addresses"`
	PublicKey [32]byte `json:"public_key"`
}

// GetPublicKey returns the public key of the peer as a BoxPublicKey
func (pi *PeerInfo) GetPublicKey() saltpack.BoxPublicKey {
	return basic.PublicKey{
		RawBoxKey: pi.PublicKey,
	}
}

// IsValid checks if the signature is valid
func (pi *PeerInfo) IsValid() bool {
	// TODO Implement
	return true
}

// NewPeerInfo from an id, an address, and a public key
func NewPeerInfo(id string, addresses []string, publicKey [32]byte) (*PeerInfo, error) {
	if id == "" {
		return nil, errors.New("missing id")
	}

	if len(addresses) == 0 {
		return nil, errors.New("missing addresses")
	}

	if len(publicKey) == 0 {
		return nil, errors.New("missing public key")
	}

	pi := &PeerInfo{
		ID:        id,
		Addresses: addresses,
		PublicKey: publicKey,
	}

	if !pi.IsValid() {
		return nil, errors.New("id and pk don't match")
	}

	return pi, nil
}
<commit_msg>Add PeerInfo Status and dates<commit_after>package mesh

import (
	"errors"
	"time"

	"github.com/keybase/saltpack"
	"github.com/keybase/saltpack/basic"
)

// PeerInfo holds the information wire needs to connect to a remote peer
type PeerInfo struct {
	ID              string    `json:"id"`
	Addresses       []string  `json:"addresses"`
	PublicKey       [32]byte  `json:"public_key"`
	Signature       []byte    `json:"signature"`
	Status          Status    `json:"status"`
	CreatedAt       time.Time `json:"create_at"`
	UpdatedAt       time.Time `json:"updated_at"`
	LastConnectedAt time.Time `json:"last_connected_at"`
}

type peerInfoClean struct {
	ID        string   `json:"id"`
	Addresses []string `json:"addresses"`
	PublicKey [32]byte `json:"public_key"`
}

// GetPublicKey returns the public key of the peer as a BoxPublicKey
func (pi *PeerInfo) GetPublicKey() saltpack.BoxPublicKey {
	return basic.PublicKey{
		RawBoxKey: pi.PublicKey,
	}
}

// IsValid checks if the signature is valid
func (pi *PeerInfo) IsValid() bool {
	// TODO Implement
	return true
}

// NewPeerInfo from an id, an address, and a public key
func NewPeerInfo(id string, addresses []string, publicKey [32]byte) (*PeerInfo, error) {
	if id == "" {
		return nil, errors.New("missing id")
	}

	if len(addresses) == 0 {
		return nil, errors.New("missing addresses")
	}

	if len(publicKey) == 0 {
		return nil, errors.New("missing public key")
	}

	pi := &PeerInfo{
		ID:        id,
		Addresses: addresses,
		PublicKey: publicKey,
		CreatedAt: time.Now(),
		UpdatedAt: time.Now(),
	}

	if !pi.IsValid() {
		return nil, errors.New("id and pk don't match")
	}

	return pi, nil
}
<|endoftext|>


<commit_before>package gumble

import (
	"encoding/binary"
	"errors"
	"io"

	"code.google.com/p/goprotobuf/proto"
	"github.com/bontibon/gumble/MumbleProto"
)

var (
	errNilMessage     = errors.New("message is nil")
	errUnknownMessage = errors.New("unknown protobuf message type")
)

type protoMessage struct {
	proto proto.Message
}

func (pm protoMessage) WriteTo(w io.Writer) (n int64, err error) {
	if pm.proto == nil {
		return 0, errNilMessage
	}
	var protoType int
	switch pm.proto.(type) {
	case *MumbleProto.Version:
		protoType = 0
	case *MumbleProto.Authenticate:
		protoType = 2
	case *MumbleProto.Ping:
		protoType = 3
	case *MumbleProto.ChannelRemove:
		protoType = 6
	case *MumbleProto.ChannelState:
		protoType = 7
	case *MumbleProto.UserRemove:
		protoType = 8
	case *MumbleProto.TextMessage:
		protoType = 11
	default:
		return 0, errUnknownMessage
	}

	if data, err := proto.Marshal(pm.proto); err != nil {
		return 0, err
	} else {
		var written int64 = 0
		wireType := uint16(protoType)
		wireLength := uint32(len(data))
		if err := binary.Write(w, binary.BigEndian, wireType); err != nil {
			return 0, err
		} else {
			written += 2
		}
		if err := binary.Write(w, binary.BigEndian, wireLength); err != nil {
			return written, err
		} else {
			written += 4
		}
		if n, err := w.Write(data); err != nil {
			return (written + int64(n)), err
		} else {
			written += int64(n)
		}
		return written, nil
	}
}

func (pm protoMessage) gumbleMessage() {
}
<commit_msg>remove return value names from protoMessage.WriteTo<commit_after>package gumble

import (
	"encoding/binary"
	"errors"
	"io"

	"code.google.com/p/goprotobuf/proto"
	"github.com/bontibon/gumble/MumbleProto"
)

var (
	errNilMessage     = errors.New("message is nil")
	errUnknownMessage = errors.New("unknown protobuf message type")
)

type protoMessage struct {
	proto proto.Message
}

func (pm protoMessage) WriteTo(w io.Writer) (int64, error) {
	if pm.proto == nil {
		return 0, errNilMessage
	}
	var protoType int
	switch pm.proto.(type) {
	case *MumbleProto.Version:
		protoType = 0
	case *MumbleProto.Authenticate:
		protoType = 2
	case *MumbleProto.Ping:
		protoType = 3
	case *MumbleProto.ChannelRemove:
		protoType = 6
	case *MumbleProto.ChannelState:
		protoType = 7
	case *MumbleProto.UserRemove:
		protoType = 8
	case *MumbleProto.TextMessage:
		protoType = 11
	default:
		return 0, errUnknownMessage
	}

	if data, err := proto.Marshal(pm.proto); err != nil {
		return 0, err
	} else {
		var written int64 = 0
		wireType := uint16(protoType)
		wireLength := uint32(len(data))
		if err := binary.Write(w, binary.BigEndian, wireType); err != nil {
			return 0, err
		} else {
			written += 2
		}
		if err := binary.Write(w, binary.BigEndian, wireLength); err != nil {
			return written, err
		} else {
			written += 4
		}
		if n, err := w.Write(data); err != nil {
			return (written + int64(n)), err
		} else {
			written += int64(n)
		}
		return written, nil
	}
}

func (pm protoMessage) gumbleMessage() {
}
<|endoftext|>