<commit_before>package migrate

import (
	"bytes"
	"database/sql"
	"errors"
	"fmt"
	"io"
	"os"
	"path"
	"regexp"
	"sort"
	"strconv"
	"strings"
	"time"

	"github.com/rubenv/sql-migrate/sqlparse"
	"gopkg.in/gorp.v1"
)

type MigrationDirection int

const (
	Up MigrationDirection = iota
	Down
)

var tableName = "gorp_migrations"
var schemaName = ""
var numberPrefixRegex = regexp.MustCompile(`^(\d+).*$`)

// TxError is returned when any error is encountered during a database
// transaction. It contains the relevant *Migration and notes it's Id in the
// Error function output.
type TxError struct {
	Migration *Migration
	Err       error
}

func newTxError(migration *PlannedMigration, err error) error {
	return &TxError{
		Migration: migration.Migration,
		Err:       err,
	}
}

func (e *TxError) Error() string {
	return e.Err.Error() + " handling " + e.Migration.Id
}

// Set the name of the table used to store migration info.
//
// Should be called before any other call such as (Exec, ExecMax, ...).
func SetTable(name string) {
	if name != "" {
		tableName = name
	}
}

// SetSchema sets the name of a schema that the migration table be referenced.
func SetSchema(name string) {
	if name != "" {
		schemaName = name
	}
}

type Migration struct {
	Id   string
	Up   []string
	Down []string

	DisableTransactionUp   bool
	DisableTransactionDown bool
}

func (m Migration) Less(other *Migration) bool {
	switch {
	case m.isNumeric() && other.isNumeric() && m.VersionInt() != other.VersionInt():
		return m.VersionInt() < other.VersionInt()
	case m.isNumeric() && !other.isNumeric():
		return true
	case !m.isNumeric() && other.isNumeric():
		return false
	default:
		return m.Id < other.Id
	}
}

func (m Migration) isNumeric() bool {
	return len(m.NumberPrefixMatches()) > 0
}

func (m Migration) NumberPrefixMatches() []string {
	return numberPrefixRegex.FindStringSubmatch(m.Id)
}

func (m Migration) VersionInt() int64 {
	v := m.NumberPrefixMatches()[1]
	value, err := strconv.ParseInt(v, 10, 64)
	if err != nil {
		panic(fmt.Sprintf("Could not parse %q into int64: %s", v, err))
	}
	return value
}

type PlannedMigration struct {
	*Migration

	DisableTransaction bool
	Queries            []string
}

type byId []*Migration

func (b byId) Len() int           { return len(b) }
func (b byId) Swap(i, j int)      { b[i], b[j] = b[j], b[i] }
func (b byId) Less(i, j int) bool { return b[i].Less(b[j]) }

type MigrationRecord struct {
	Id        string    `db:"id"`
	AppliedAt time.Time `db:"applied_at"`
}

var MigrationDialects = map[string]gorp.Dialect{
	"sqlite3":  gorp.SqliteDialect{},
	"postgres": gorp.PostgresDialect{},
	"mysql":    gorp.MySQLDialect{Engine: "InnoDB", Encoding: "UTF8"},
	"mssql":    gorp.SqlServerDialect{},
	"oci8":     gorp.OracleDialect{},
}

type MigrationSource interface {
	// Finds the migrations.
	//
	// The resulting slice of migrations should be sorted by Id.
	FindMigrations() ([]*Migration, error)
}

// A hardcoded set of migrations, in-memory.
type MemoryMigrationSource struct {
	Migrations []*Migration
}

var _ MigrationSource = (*MemoryMigrationSource)(nil)

func (m MemoryMigrationSource) FindMigrations() ([]*Migration, error) {
	// Make sure migrations are sorted
	sort.Sort(byId(m.Migrations))

	return m.Migrations, nil
}

// A set of migrations loaded from a directory.
type FileMigrationSource struct {
	Dir string
}

var _ MigrationSource = (*FileMigrationSource)(nil)

func (f FileMigrationSource) FindMigrations() ([]*Migration, error) {
	migrations := make([]*Migration, 0)

	file, err := os.Open(f.Dir)
	if err != nil {
		return nil, err
	}

	files, err := file.Readdir(0)
	if err != nil {
		return nil, err
	}

	for _, info := range files {
		if strings.HasSuffix(info.Name(), ".sql") {
			file, err := os.Open(path.Join(f.Dir, info.Name()))
			if err != nil {
				return nil, err
			}

			migration, err := ParseMigration(info.Name(), file)
			if err != nil {
				return nil, err
			}

			migrations = append(migrations, migration)
		}
	}

	// Make sure migrations are sorted
	sort.Sort(byId(migrations))

	return migrations, nil
}

// Migrations from a bindata asset set.
type AssetMigrationSource struct {
	// Asset should return content of file in path if exists
	Asset func(path string) ([]byte, error)

	// AssetDir should return list of files in the path
	AssetDir func(path string) ([]string, error)

	// Path in the bindata to use.
	Dir string
}

var _ MigrationSource = (*AssetMigrationSource)(nil)

func (a AssetMigrationSource) FindMigrations() ([]*Migration, error) {
	migrations := make([]*Migration, 0)

	files, err := a.AssetDir(a.Dir)
	if err != nil {
		return nil, err
	}

	for _, name := range files {
		if strings.HasSuffix(name, ".sql") {
			file, err := a.Asset(path.Join(a.Dir, name))
			if err != nil {
				return nil, err
			}

			migration, err := ParseMigration(name, bytes.NewReader(file))
			if err != nil {
				return nil, err
			}

			migrations = append(migrations, migration)
		}
	}

	// Make sure migrations are sorted
	sort.Sort(byId(migrations))

	return migrations, nil
}

// Migration parsing
func ParseMigration(id string, r io.ReadSeeker) (*Migration, error) {
	m := &Migration{
		Id: id,
	}

	parsed, err := sqlparse.ParseMigration(r)
	if err != nil {
		return nil, err
	}

	m.Up = parsed.UpStatements
	m.Down = parsed.DownStatements

	m.DisableTransactionUp = parsed.DisableTransactionUp
	m.DisableTransactionDown = parsed.DisableTransactionDown

	return m, nil
}

type SqlExecutor interface {
	Exec(query string, args ...interface{}) (sql.Result, error)
	Insert(list ...interface{}) error
	Delete(list ...interface{}) (int64, error)
}

// Execute a set of migrations
//
// Returns the number of applied migrations.
func Exec(db *sql.DB, dialect string, m MigrationSource, dir MigrationDirection) (int, error) {
	return ExecMax(db, dialect, m, dir, 0)
}

// Execute a set of migrations
//
// Will apply at most `max` migrations. Pass 0 for no limit (or use Exec).
//
// Returns the number of applied migrations.
func ExecMax(db *sql.DB, dialect string, m MigrationSource, dir MigrationDirection, max int) (int, error) {
	migrations, dbMap, err := PlanMigration(db, dialect, m, dir, max)
	if err != nil {
		return 0, err
	}

	// Apply migrations
	applied := 0
	for _, migration := range migrations {
		var executor SqlExecutor

		if migration.DisableTransaction {
			executor = dbMap
		} else {
			executor, err = dbMap.Begin()
			if err != nil {
				return applied, newTxError(migration, err)
			}
		}

		for _, stmt := range migration.Queries {
			if _, err := executor.Exec(stmt); err != nil {
				if trans, ok := executor.(*gorp.Transaction); ok {
					trans.Rollback()
				}

				return applied, newTxError(migration, err)
			}
		}

		if dir == Up {
			err = executor.Insert(&MigrationRecord{
				Id:        migration.Id,
				AppliedAt: time.Now(),
			})
			if err != nil {
				return applied, newTxError(migration, err)
			}
		} else if dir == Down {
			_, err := executor.Delete(&MigrationRecord{
				Id: migration.Id,
			})
			if err != nil {
				return applied, newTxError(migration, err)
			}
		} else {
			panic("Not possible")
		}

		if trans, ok := executor.(*gorp.Transaction); ok {
			if err := trans.Commit(); err != nil {
				return applied, newTxError(migration, err)
			}
		}

		applied++
	}

	return applied, nil
}

// Plan a migration.
func PlanMigration(db *sql.DB, dialect string, m MigrationSource, dir MigrationDirection, max int) ([]*PlannedMigration, *gorp.DbMap, error) {
	dbMap, err := getMigrationDbMap(db, dialect)
	if err != nil {
		return nil, nil, err
	}

	migrations, err := m.FindMigrations()
	if err != nil {
		return nil, nil, err
	}

	var migrationRecords []MigrationRecord
	_, err = dbMap.Select(&migrationRecords, fmt.Sprintf("SELECT * FROM %s", dbMap.Dialect.QuotedTableForQuery(schemaName, tableName)))
	if err != nil {
		return nil, nil, err
	}

	// Sort migrations that have been run by Id.
	var existingMigrations []*Migration
	for _, migrationRecord := range migrationRecords {
		existingMigrations = append(existingMigrations, &Migration{
			Id: migrationRecord.Id,
		})
	}
	sort.Sort(byId(existingMigrations))

	// Get last migration that was run
	record := &Migration{}
	if len(existingMigrations) > 0 {
		record = existingMigrations[len(existingMigrations)-1]
	}

	result := make([]*PlannedMigration, 0)

	// Add missing migrations up to the last run migration.
	// This can happen for example when merges happened.
	if len(existingMigrations) > 0 {
		result = append(result, ToCatchup(migrations, existingMigrations, record)...)
	}

	// Figure out which migrations to apply
	toApply := ToApply(migrations, record.Id, dir)
	toApplyCount := len(toApply)
	if max > 0 && max < toApplyCount {
		toApplyCount = max
	}
	for _, v := range toApply[0:toApplyCount] {

		if dir == Up {
			result = append(result, &PlannedMigration{
				Migration:          v,
				Queries:            v.Up,
				DisableTransaction: v.DisableTransactionUp,
			})
		} else if dir == Down {
			result = append(result, &PlannedMigration{
				Migration:          v,
				Queries:            v.Down,
				DisableTransaction: v.DisableTransactionDown,
			})
		}
	}

	return result, dbMap, nil
}

// Filter a slice of migrations into ones that should be applied.
func ToApply(migrations []*Migration, current string, direction MigrationDirection) []*Migration {
	var index = -1
	if current != "" {
		for index < len(migrations)-1 {
			index++
			if migrations[index].Id == current {
				break
			}
		}
	}

	if direction == Up {
		return migrations[index+1:]
	} else if direction == Down {
		if index == -1 {
			return []*Migration{}
		}

		// Add in reverse order
		toApply := make([]*Migration, index+1)
		for i := 0; i < index+1; i++ {
			toApply[index-i] = migrations[i]
		}
		return toApply
	}

	panic("Not possible")
}

func ToCatchup(migrations, existingMigrations []*Migration, lastRun *Migration) []*PlannedMigration {
	missing := make([]*PlannedMigration, 0)
	for _, migration := range migrations {
		found := false
		for _, existing := range existingMigrations {
			if existing.Id == migration.Id {
				found = true
				break
			}
		}
		if !found && migration.Less(lastRun) {
			missing = append(missing, &PlannedMigration{
				Migration:          migration,
				Queries:            migration.Up,
				DisableTransaction: migration.DisableTransactionUp,
			})
		}
	}
	return missing
}

func GetMigrationRecords(db *sql.DB, dialect string) ([]*MigrationRecord, error) {
	dbMap, err := getMigrationDbMap(db, dialect)
	if err != nil {
		return nil, err
	}

	var records []*MigrationRecord
	query := fmt.Sprintf("SELECT * FROM %s ORDER BY id ASC", dbMap.Dialect.QuotedTableForQuery(schemaName, tableName))
	_, err = dbMap.Select(&records, query)
	if err != nil {
		return nil, err
	}

	return records, nil
}

func getMigrationDbMap(db *sql.DB, dialect string) (*gorp.DbMap, error) {
	d, ok := MigrationDialects[dialect]
	if !ok {
		return nil, fmt.Errorf("Unknown dialect: %s", dialect)
	}

	// When using the mysql driver, make sure that the parseTime option is
	// configured, otherwise it won't map time columns to time.Time. See
	// https://github.com/rubenv/sql-migrate/issues/2
	if dialect == "mysql" {
		var out *time.Time
		err := db.QueryRow("SELECT NOW()").Scan(&out)
		if err != nil {
			if err.Error() == "sql: Scan error on column index 0: unsupported driver -> Scan pair: []uint8 -> *time.Time" ||
				err.Error() == "sql: Scan error on column index 0: unsupported Scan, storing driver.Value type []uint8 into type *time.Time" {
				return nil, errors.New(`Cannot parse dates.

Make sure that the parseTime option is supplied to your database connection.
Check https://github.com/go-sql-driver/mysql#parsetime for more info.`)
			} else {
				return nil, err
			}
		}
	}

	// Create migration database map
	dbMap := &gorp.DbMap{Db: db, Dialect: d}
	dbMap.AddTableWithNameAndSchema(MigrationRecord{}, schemaName, tableName).SetKeys(false, "Id")
	//dbMap.TraceOn("", log.New(os.Stdout, "migrate: ", log.Lmicroseconds))

	err := dbMap.CreateTablesIfNotExists()
	if err != nil {
		return nil, err
	}

	return dbMap, nil
}

// TODO: Run migration + record insert in transaction.
<commit_msg>feat: added more detailed error message upon which migration is actually invalid while discovering/parsing them<commit_after>package migrate

import (
	"bytes"
	"database/sql"
	"errors"
	"fmt"
	"io"
	"os"
	"path"
	"regexp"
	"sort"
	"strconv"
	"strings"
	"time"

	"github.com/rubenv/sql-migrate/sqlparse"
	"gopkg.in/gorp.v1"
)

type MigrationDirection int

const (
	Up MigrationDirection = iota
	Down
)

var tableName = "gorp_migrations"
var schemaName = ""
var numberPrefixRegex = regexp.MustCompile(`^(\d+).*$`)

// TxError is returned when any error is encountered during a database
// transaction. It contains the relevant *Migration and notes it's Id in the
// Error function output.
type TxError struct {
	Migration *Migration
	Err       error
}

func newTxError(migration *PlannedMigration, err error) error {
	return &TxError{
		Migration: migration.Migration,
		Err:       err,
	}
}

func (e *TxError) Error() string {
	return e.Err.Error() + " handling " + e.Migration.Id
}

// Set the name of the table used to store migration info.
//
// Should be called before any other call such as (Exec, ExecMax, ...).
func SetTable(name string) {
	if name != "" {
		tableName = name
	}
}

// SetSchema sets the name of a schema that the migration table be referenced.
func SetSchema(name string) {
	if name != "" {
		schemaName = name
	}
}

type Migration struct {
	Id   string
	Up   []string
	Down []string

	DisableTransactionUp   bool
	DisableTransactionDown bool
}

func (m Migration) Less(other *Migration) bool {
	switch {
	case m.isNumeric() && other.isNumeric() && m.VersionInt() != other.VersionInt():
		return m.VersionInt() < other.VersionInt()
	case m.isNumeric() && !other.isNumeric():
		return true
	case !m.isNumeric() && other.isNumeric():
		return false
	default:
		return m.Id < other.Id
	}
}

func (m Migration) isNumeric() bool {
	return len(m.NumberPrefixMatches()) > 0
}

func (m Migration) NumberPrefixMatches() []string {
	return numberPrefixRegex.FindStringSubmatch(m.Id)
}

func (m Migration) VersionInt() int64 {
	v := m.NumberPrefixMatches()[1]
	value, err := strconv.ParseInt(v, 10, 64)
	if err != nil {
		panic(fmt.Sprintf("Could not parse %q into int64: %s", v, err))
	}
	return value
}

type PlannedMigration struct {
	*Migration

	DisableTransaction bool
	Queries            []string
}

type byId []*Migration

func (b byId) Len() int           { return len(b) }
func (b byId) Swap(i, j int)      { b[i], b[j] = b[j], b[i] }
func (b byId) Less(i, j int) bool { return b[i].Less(b[j]) }

type MigrationRecord struct {
	Id        string    `db:"id"`
	AppliedAt time.Time `db:"applied_at"`
}

var MigrationDialects = map[string]gorp.Dialect{
	"sqlite3":  gorp.SqliteDialect{},
	"postgres": gorp.PostgresDialect{},
	"mysql":    gorp.MySQLDialect{Engine: "InnoDB", Encoding: "UTF8"},
	"mssql":    gorp.SqlServerDialect{},
	"oci8":     gorp.OracleDialect{},
}

type MigrationSource interface {
	// Finds the migrations.
	//
	// The resulting slice of migrations should be sorted by Id.
	FindMigrations() ([]*Migration, error)
}

// A hardcoded set of migrations, in-memory.
type MemoryMigrationSource struct {
	Migrations []*Migration
}

var _ MigrationSource = (*MemoryMigrationSource)(nil)

func (m MemoryMigrationSource) FindMigrations() ([]*Migration, error) {
	// Make sure migrations are sorted
	sort.Sort(byId(m.Migrations))

	return m.Migrations, nil
}

// A set of migrations loaded from a directory.
type FileMigrationSource struct {
	Dir string
}

var _ MigrationSource = (*FileMigrationSource)(nil)

func (f FileMigrationSource) FindMigrations() ([]*Migration, error) {
	migrations := make([]*Migration, 0)

	file, err := os.Open(f.Dir)
	if err != nil {
		return nil, err
	}

	files, err := file.Readdir(0)
	if err != nil {
		return nil, err
	}

	for _, info := range files {
		if strings.HasSuffix(info.Name(), ".sql") {
			file, err := os.Open(path.Join(f.Dir, info.Name()))
			if err != nil {
				return nil, fmt.Errorf("Error while opening %s : %s", info.Name(), err)
			}

			migration, err := ParseMigration(info.Name(), file)
			if err != nil {
				return nil, fmt.Errorf("Error while parsing %s : %s", info.Name(), err)
			}

			migrations = append(migrations, migration)
		}
	}

	// Make sure migrations are sorted
	sort.Sort(byId(migrations))

	return migrations, nil
}

// Migrations from a bindata asset set.
type AssetMigrationSource struct {
	// Asset should return content of file in path if exists
	Asset func(path string) ([]byte, error)

	// AssetDir should return list of files in the path
	AssetDir func(path string) ([]string, error)

	// Path in the bindata to use.
	Dir string
}

var _ MigrationSource = (*AssetMigrationSource)(nil)

func (a AssetMigrationSource) FindMigrations() ([]*Migration, error) {
	migrations := make([]*Migration, 0)

	files, err := a.AssetDir(a.Dir)
	if err != nil {
		return nil, err
	}

	for _, name := range files {
		if strings.HasSuffix(name, ".sql") {
			file, err := a.Asset(path.Join(a.Dir, name))
			if err != nil {
				return nil, err
			}

			migration, err := ParseMigration(name, bytes.NewReader(file))
			if err != nil {
				return nil, err
			}

			migrations = append(migrations, migration)
		}
	}

	// Make sure migrations are sorted
	sort.Sort(byId(migrations))

	return migrations, nil
}

// Migration parsing
func ParseMigration(id string, r io.ReadSeeker) (*Migration, error) {
	m := &Migration{
		Id: id,
	}

	parsed, err := sqlparse.ParseMigration(r)
	if err != nil {
		return nil, err
	}

	m.Up = parsed.UpStatements
	m.Down = parsed.DownStatements

	m.DisableTransactionUp = parsed.DisableTransactionUp
	m.DisableTransactionDown = parsed.DisableTransactionDown

	return m, nil
}

type SqlExecutor interface {
	Exec(query string, args ...interface{}) (sql.Result, error)
	Insert(list ...interface{}) error
	Delete(list ...interface{}) (int64, error)
}

// Execute a set of migrations
//
// Returns the number of applied migrations.
func Exec(db *sql.DB, dialect string, m MigrationSource, dir MigrationDirection) (int, error) {
	return ExecMax(db, dialect, m, dir, 0)
}

// Execute a set of migrations
//
// Will apply at most `max` migrations. Pass 0 for no limit (or use Exec).
//
// Returns the number of applied migrations.
func ExecMax(db *sql.DB, dialect string, m MigrationSource, dir MigrationDirection, max int) (int, error) {
	migrations, dbMap, err := PlanMigration(db, dialect, m, dir, max)
	if err != nil {
		return 0, err
	}

	// Apply migrations
	applied := 0
	for _, migration := range migrations {
		var executor SqlExecutor

		if migration.DisableTransaction {
			executor = dbMap
		} else {
			executor, err = dbMap.Begin()
			if err != nil {
				return applied, newTxError(migration, err)
			}
		}

		for _, stmt := range migration.Queries {
			if _, err := executor.Exec(stmt); err != nil {
				if trans, ok := executor.(*gorp.Transaction); ok {
					trans.Rollback()
				}

				return applied, newTxError(migration, err)
			}
		}

		if dir == Up {
			err = executor.Insert(&MigrationRecord{
				Id:        migration.Id,
				AppliedAt: time.Now(),
			})
			if err != nil {
				return applied, newTxError(migration, err)
			}
		} else if dir == Down {
			_, err := executor.Delete(&MigrationRecord{
				Id: migration.Id,
			})
			if err != nil {
				return applied, newTxError(migration, err)
			}
		} else {
			panic("Not possible")
		}

		if trans, ok := executor.(*gorp.Transaction); ok {
			if err := trans.Commit(); err != nil {
				return applied, newTxError(migration, err)
			}
		}

		applied++
	}

	return applied, nil
}

// Plan a migration.
func PlanMigration(db *sql.DB, dialect string, m MigrationSource, dir MigrationDirection, max int) ([]*PlannedMigration, *gorp.DbMap, error) {
	dbMap, err := getMigrationDbMap(db, dialect)
	if err != nil {
		return nil, nil, err
	}

	migrations, err := m.FindMigrations()
	if err != nil {
		return nil, nil, err
	}

	var migrationRecords []MigrationRecord
	_, err = dbMap.Select(&migrationRecords, fmt.Sprintf("SELECT * FROM %s", dbMap.Dialect.QuotedTableForQuery(schemaName, tableName)))
	if err != nil {
		return nil, nil, err
	}

	// Sort migrations that have been run by Id.
	var existingMigrations []*Migration
	for _, migrationRecord := range migrationRecords {
		existingMigrations = append(existingMigrations, &Migration{
			Id: migrationRecord.Id,
		})
	}
	sort.Sort(byId(existingMigrations))

	// Get last migration that was run
	record := &Migration{}
	if len(existingMigrations) > 0 {
		record = existingMigrations[len(existingMigrations)-1]
	}

	result := make([]*PlannedMigration, 0)

	// Add missing migrations up to the last run migration.
	// This can happen for example when merges happened.
	if len(existingMigrations) > 0 {
		result = append(result, ToCatchup(migrations, existingMigrations, record)...)
	}

	// Figure out which migrations to apply
	toApply := ToApply(migrations, record.Id, dir)
	toApplyCount := len(toApply)
	if max > 0 && max < toApplyCount {
		toApplyCount = max
	}
	for _, v := range toApply[0:toApplyCount] {

		if dir == Up {
			result = append(result, &PlannedMigration{
				Migration:          v,
				Queries:            v.Up,
				DisableTransaction: v.DisableTransactionUp,
			})
		} else if dir == Down {
			result = append(result, &PlannedMigration{
				Migration:          v,
				Queries:            v.Down,
				DisableTransaction: v.DisableTransactionDown,
			})
		}
	}

	return result, dbMap, nil
}

// Filter a slice of migrations into ones that should be applied.
func ToApply(migrations []*Migration, current string, direction MigrationDirection) []*Migration {
	var index = -1
	if current != "" {
		for index < len(migrations)-1 {
			index++
			if migrations[index].Id == current {
				break
			}
		}
	}

	if direction == Up {
		return migrations[index+1:]
	} else if direction == Down {
		if index == -1 {
			return []*Migration{}
		}

		// Add in reverse order
		toApply := make([]*Migration, index+1)
		for i := 0; i < index+1; i++ {
			toApply[index-i] = migrations[i]
		}
		return toApply
	}

	panic("Not possible")
}

func ToCatchup(migrations, existingMigrations []*Migration, lastRun *Migration) []*PlannedMigration {
	missing := make([]*PlannedMigration, 0)
	for _, migration := range migrations {
		found := false
		for _, existing := range existingMigrations {
			if existing.Id == migration.Id {
				found = true
				break
			}
		}
		if !found && migration.Less(lastRun) {
			missing = append(missing, &PlannedMigration{
				Migration:          migration,
				Queries:            migration.Up,
				DisableTransaction: migration.DisableTransactionUp,
			})
		}
	}
	return missing
}

func GetMigrationRecords(db *sql.DB, dialect string) ([]*MigrationRecord, error) {
	dbMap, err := getMigrationDbMap(db, dialect)
	if err != nil {
		return nil, err
	}

	var records []*MigrationRecord
	query := fmt.Sprintf("SELECT * FROM %s ORDER BY id ASC", dbMap.Dialect.QuotedTableForQuery(schemaName, tableName))
	_, err = dbMap.Select(&records, query)
	if err != nil {
		return nil, err
	}

	return records, nil
}

func getMigrationDbMap(db *sql.DB, dialect string) (*gorp.DbMap, error) {
	d, ok := MigrationDialects[dialect]
	if !ok {
		return nil, fmt.Errorf("Unknown dialect: %s", dialect)
	}

	// When using the mysql driver, make sure that the parseTime option is
	// configured, otherwise it won't map time columns to time.Time. See
	// https://github.com/rubenv/sql-migrate/issues/2
	if dialect == "mysql" {
		var out *time.Time
		err := db.QueryRow("SELECT NOW()").Scan(&out)
		if err != nil {
			if err.Error() == "sql: Scan error on column index 0: unsupported driver -> Scan pair: []uint8 -> *time.Time" ||
				err.Error() == "sql: Scan error on column index 0: unsupported Scan, storing driver.Value type []uint8 into type *time.Time" {
				return nil, errors.New(`Cannot parse dates.

Make sure that the parseTime option is supplied to your database connection.
Check https://github.com/go-sql-driver/mysql#parsetime for more info.`)
			} else {
				return nil, err
			}
		}
	}

	// Create migration database map
	dbMap := &gorp.DbMap{Db: db, Dialect: d}
	dbMap.AddTableWithNameAndSchema(MigrationRecord{}, schemaName, tableName).SetKeys(false, "Id")
	//dbMap.TraceOn("", log.New(os.Stdout, "migrate: ", log.Lmicroseconds))

	err := dbMap.CreateTablesIfNotExists()
	if err != nil {
		return nil, err
	}

	return dbMap, nil
}

// TODO: Run migration + record insert in transaction.
<|endoftext|>


<commit_before>package physical

import (
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"sync"
)

// FileBackend is a physical backend that stores data on disk
// at a given file path. It can be used for durable single server
// situations, or to develop locally where durability is not critical.
//
// WARNING: the file backend implementation is currently extremely unsafe
// and non-performant. It is meant mostly for local testing and development.
// It can be improved in the future.
type FileBackend struct {
	Path string

	l sync.Mutex
}

// newFileBackend constructs a Filebackend using the given directory
func newFileBackend(conf map[string]string) (Backend, error) {
	path, ok := conf["path"]
	if !ok {
		return nil, fmt.Errorf("'path' must be set")
	}

	return &FileBackend{Path: path}, nil
}

func (b *FileBackend) Delete(k string) error {
	b.l.Lock()
	defer b.l.Unlock()

	path, key := b.path(k)
	path = filepath.Join(path, key)

	err := os.Remove(path)
	if err != nil && os.IsNotExist(err) {
		err = nil
	}

	return err
}

func (b *FileBackend) Get(k string) (*Entry, error) {
	b.l.Lock()
	defer b.l.Unlock()

	path, key := b.path(k)
	path = filepath.Join(path, key)

	f, err := os.Open(path)
	if err != nil {
		if os.IsNotExist(err) {
			return nil, nil
		}

		return nil, err
	}
	defer f.Close()

	var entry Entry
	dec := json.NewDecoder(f)
	if err := dec.Decode(&entry); err != nil {
		return nil, err
	}

	return &entry, nil
}

func (b *FileBackend) Put(entry *Entry) error {
	path, key := b.path(entry.Key)

	b.l.Lock()
	defer b.l.Unlock()

	// Make the parent tree
	if err := os.MkdirAll(path, 0755); err != nil {
		return err
	}

	// JSON encode the entry and write it
	f, err := os.OpenFile(filepath.Join(path, key), os.O_CREATE|os.O_TRUNC, 0600)
	if err != nil {
		return err
	}
	defer f.Close()
	enc := json.NewEncoder(f)
	return enc.Encode(entry)
}

func (b *FileBackend) List(prefix string) ([]string, error) {
	b.l.Lock()
	defer b.l.Unlock()

	path := b.Path
	if prefix != "" {
		path = filepath.Join(path, prefix)
	}

	// Read the directory contents
	f, err := os.Open(path)
	if err != nil {
		if os.IsNotExist(err) {
			return nil, nil
		}

		return nil, err
	}
	defer f.Close()

	names, err := f.Readdirnames(-1)
	if err != nil {
		return nil, err
	}

	for i, name := range names {
		if name[0] == '_' {
			names[i] = name[1:]
		} else {
			names[i] = name + "/"
		}
	}

	return names, nil
}

func (b *FileBackend) path(k string) (string, string) {
	path := filepath.Join(b.Path, k)
	key := filepath.Base(path)
	path = filepath.Dir(path)
	return path, "_" + key
}
<commit_msg>physical/file: open for writing<commit_after>package physical

import (
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"sync"
)

// FileBackend is a physical backend that stores data on disk
// at a given file path. It can be used for durable single server
// situations, or to develop locally where durability is not critical.
//
// WARNING: the file backend implementation is currently extremely unsafe
// and non-performant. It is meant mostly for local testing and development.
// It can be improved in the future.
type FileBackend struct {
	Path string

	l sync.Mutex
}

// newFileBackend constructs a Filebackend using the given directory
func newFileBackend(conf map[string]string) (Backend, error) {
	path, ok := conf["path"]
	if !ok {
		return nil, fmt.Errorf("'path' must be set")
	}

	return &FileBackend{Path: path}, nil
}

func (b *FileBackend) Delete(k string) error {
	b.l.Lock()
	defer b.l.Unlock()

	path, key := b.path(k)
	path = filepath.Join(path, key)

	err := os.Remove(path)
	if err != nil && os.IsNotExist(err) {
		err = nil
	}

	return err
}

func (b *FileBackend) Get(k string) (*Entry, error) {
	b.l.Lock()
	defer b.l.Unlock()

	path, key := b.path(k)
	path = filepath.Join(path, key)

	f, err := os.Open(path)
	if err != nil {
		if os.IsNotExist(err) {
			return nil, nil
		}

		return nil, err
	}
	defer f.Close()

	var entry Entry
	dec := json.NewDecoder(f)
	if err := dec.Decode(&entry); err != nil {
		return nil, err
	}

	return &entry, nil
}

func (b *FileBackend) Put(entry *Entry) error {
	path, key := b.path(entry.Key)

	b.l.Lock()
	defer b.l.Unlock()

	// Make the parent tree
	if err := os.MkdirAll(path, 0755); err != nil {
		return err
	}

	// JSON encode the entry and write it
	f, err := os.OpenFile(
		filepath.Join(path, key),
		os.O_CREATE|os.O_TRUNC|os.O_WRONLY,
		0600)
	if err != nil {
		return err
	}
	defer f.Close()
	enc := json.NewEncoder(f)
	return enc.Encode(entry)
}

func (b *FileBackend) List(prefix string) ([]string, error) {
	b.l.Lock()
	defer b.l.Unlock()

	path := b.Path
	if prefix != "" {
		path = filepath.Join(path, prefix)
	}

	// Read the directory contents
	f, err := os.Open(path)
	if err != nil {
		if os.IsNotExist(err) {
			return nil, nil
		}

		return nil, err
	}
	defer f.Close()

	names, err := f.Readdirnames(-1)
	if err != nil {
		return nil, err
	}

	for i, name := range names {
		if name[0] == '_' {
			names[i] = name[1:]
		} else {
			names[i] = name + "/"
		}
	}

	return names, nil
}

func (b *FileBackend) path(k string) (string, string) {
	path := filepath.Join(b.Path, k)
	key := filepath.Base(path)
	path = filepath.Dir(path)
	return path, "_" + key
}
<|endoftext|>


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

// This file contains tests verifying the types associated with an AST after
// type checking.

package types

import (
	"go/ast"
	"go/parser"
	"testing"
)

const filename = "<src>"

func makePkg(t *testing.T, src string) (*ast.Package, error) {
	file, err := parser.ParseFile(fset, filename, src, parser.DeclarationErrors)
	if err != nil {
		return nil, err
	}
	files := map[string]*ast.File{filename: file}
	pkg, err := ast.NewPackage(fset, files, GcImport, Universe)
	if err != nil {
		return nil, err
	}
	if err := Check(fset, pkg, nil, nil); err != nil {
		return nil, err
	}
	return pkg, nil
}

type testEntry struct {
	src, str string
}

// dup returns a testEntry where both src and str are the same.
func dup(s string) testEntry {
	return testEntry{s, s}
}

var testTypes = []testEntry{
	// basic types
	dup("int"),
	dup("float32"),
	dup("string"),

	// arrays
	dup("[10]int"),

	// slices
	dup("[]int"),
	dup("[][]int"),

	// structs
	dup("struct{}"),
	dup("struct{x int}"),
	{`struct {
		x, y int
		z float32 "foo"
	}`, `struct{x int; y int; z float32 "foo"}`},
	{`struct {
		string
		elems []T
	}`, `struct{string; elems []T}`},

	// pointers
	dup("*int"),
	dup("***struct{}"),
	dup("*struct{a int; b float32}"),

	// functions
	dup("func()"),
	dup("func(x int)"),
	{"func(x, y int)", "func(x int, y int)"},
	{"func(x, y int, z string)", "func(x int, y int, z string)"},
	dup("func(int)"),
	{"func(int, string, byte)", "func(int, string, byte)"},

	dup("func() int"),
	{"func() (string)", "func() string"},
	dup("func() (u int)"),
	{"func() (u, v int, w string)", "func() (u int, v int, w string)"},

	dup("func(int) string"),
	dup("func(x int) string"),
	dup("func(x int) (u string)"),
	{"func(x, y int) (u string)", "func(x int, y int) (u string)"},

	dup("func(...int) string"),
	dup("func(x ...int) string"),
	dup("func(x ...int) (u string)"),
	{"func(x, y ...int) (u string)", "func(x int, y ...int) (u string)"},

	// interfaces
	dup("interface{}"),
	dup("interface{m()}"),
	{`interface{
		m(int) float32
		String() string
	}`, `interface{String() string; m(int) float32}`}, // methods are sorted
	// TODO(gri) add test for interface w/ anonymous field

	// maps
	dup("map[string]int"),
	{"map[struct{x, y int}][]byte", "map[struct{x int; y int}][]byte"},

	// channels
	dup("chan int"),
	dup("chan<- func()"),
	dup("<-chan []func() int"),
}

func TestTypes(t *testing.T) {
	for _, test := range testTypes {
		src := "package p; type T " + test.src
		pkg, err := makePkg(t, src)
		if err != nil {
			t.Errorf("%s: %s", src, err)
			continue
		}
		typ := underlying(pkg.Scope.Lookup("T").Type.(Type))
		str := typeString(typ)
		if str != test.str {
			t.Errorf("%s: got %s, want %s", test.src, str, test.str)
		}
	}
}

var testExprs = []testEntry{
	// basic type literals
	dup("x"),
	dup("true"),
	dup("42"),
	dup("3.1415"),
	dup("2.71828i"),
	dup(`'a'`),
	dup(`"foo"`),
	dup("`bar`"),

	// arbitrary expressions
	dup("&x"),
	dup("*x"),
	dup("(x)"),
	dup("x + y"),
	dup("x + y * 10"),
	dup("s.foo"),
	dup("s[0]"),
	dup("s[x:y]"),
	dup("s[:y]"),
	dup("s[x:]"),
	dup("s[:]"),
	dup("f(1, 2.3)"),
	dup("-f(10, 20)"),
	dup("f(x + y, +3.1415)"),
	{"func(a, b int) {}", "(func literal)"},
	{"func(a, b int) []int {}()[x]", "(func literal)()[x]"},
	{"[]int{1, 2, 3}", "(composite literal)"},
	{"[]int{1, 2, 3}[x:]", "(composite literal)[x:]"},
	{"x.([]string)", "x.(...)"},
}

func TestExprs(t *testing.T) {
	for _, test := range testExprs {
		src := "package p; var _ = " + test.src + "; var (x, y int; s []string; f func(int, float32))"
		pkg, err := makePkg(t, src)
		if err != nil {
			t.Errorf("%s: %s", src, err)
			continue
		}
		expr := pkg.Files[filename].Decls[0].(*ast.GenDecl).Specs[0].(*ast.ValueSpec).Values[0]
		str := exprString(expr)
		if str != test.str {
			t.Errorf("%s: got %s, want %s", test.src, str, test.str)
		}
	}
}
<commit_msg>fix build: use temporary variable to avoid compiler error<commit_after>// Copyright 2012 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// This file contains tests verifying the types associated with an AST after
// type checking.

package types

import (
	"go/ast"
	"go/parser"
	"testing"
)

const filename = "<src>"

func makePkg(t *testing.T, src string) (*ast.Package, error) {
	file, err := parser.ParseFile(fset, filename, src, parser.DeclarationErrors)
	if err != nil {
		return nil, err
	}
	files := map[string]*ast.File{filename: file}
	pkg, err := ast.NewPackage(fset, files, GcImport, Universe)
	if err != nil {
		return nil, err
	}
	if err := Check(fset, pkg, nil, nil); err != nil {
		return nil, err
	}
	return pkg, nil
}

type testEntry struct {
	src, str string
}

// dup returns a testEntry where both src and str are the same.
func dup(s string) testEntry {
	return testEntry{s, s}
}

var testTypes = []testEntry{
	// basic types
	dup("int"),
	dup("float32"),
	dup("string"),

	// arrays
	dup("[10]int"),

	// slices
	dup("[]int"),
	dup("[][]int"),

	// structs
	dup("struct{}"),
	dup("struct{x int}"),
	{`struct {
		x, y int
		z float32 "foo"
	}`, `struct{x int; y int; z float32 "foo"}`},
	{`struct {
		string
		elems []T
	}`, `struct{string; elems []T}`},

	// pointers
	dup("*int"),
	dup("***struct{}"),
	dup("*struct{a int; b float32}"),

	// functions
	dup("func()"),
	dup("func(x int)"),
	{"func(x, y int)", "func(x int, y int)"},
	{"func(x, y int, z string)", "func(x int, y int, z string)"},
	dup("func(int)"),
	{"func(int, string, byte)", "func(int, string, byte)"},

	dup("func() int"),
	{"func() (string)", "func() string"},
	dup("func() (u int)"),
	{"func() (u, v int, w string)", "func() (u int, v int, w string)"},

	dup("func(int) string"),
	dup("func(x int) string"),
	dup("func(x int) (u string)"),
	{"func(x, y int) (u string)", "func(x int, y int) (u string)"},

	dup("func(...int) string"),
	dup("func(x ...int) string"),
	dup("func(x ...int) (u string)"),
	{"func(x, y ...int) (u string)", "func(x int, y ...int) (u string)"},

	// interfaces
	dup("interface{}"),
	dup("interface{m()}"),
	{`interface{
		m(int) float32
		String() string
	}`, `interface{String() string; m(int) float32}`}, // methods are sorted
	// TODO(gri) add test for interface w/ anonymous field

	// maps
	dup("map[string]int"),
	{"map[struct{x, y int}][]byte", "map[struct{x int; y int}][]byte"},

	// channels
	dup("chan int"),
	dup("chan<- func()"),
	dup("<-chan []func() int"),
}

func TestTypes(t *testing.T) {
	for _, test := range testTypes {
		src := "package p; type T " + test.src
		pkg, err := makePkg(t, src)
		if err != nil {
			t.Errorf("%s: %s", src, err)
			continue
		}
		typ := underlying(pkg.Scope.Lookup("T").Type.(Type))
		str := typeString(typ)
		if str != test.str {
			t.Errorf("%s: got %s, want %s", test.src, str, test.str)
		}
	}
}

var testExprs = []testEntry{
	// basic type literals
	dup("x"),
	dup("true"),
	dup("42"),
	dup("3.1415"),
	dup("2.71828i"),
	dup(`'a'`),
	dup(`"foo"`),
	dup("`bar`"),

	// arbitrary expressions
	dup("&x"),
	dup("*x"),
	dup("(x)"),
	dup("x + y"),
	dup("x + y * 10"),
	dup("s.foo"),
	dup("s[0]"),
	dup("s[x:y]"),
	dup("s[:y]"),
	dup("s[x:]"),
	dup("s[:]"),
	dup("f(1, 2.3)"),
	dup("-f(10, 20)"),
	dup("f(x + y, +3.1415)"),
	{"func(a, b int) {}", "(func literal)"},
	{"func(a, b int) []int {}()[x]", "(func literal)()[x]"},
	{"[]int{1, 2, 3}", "(composite literal)"},
	{"[]int{1, 2, 3}[x:]", "(composite literal)[x:]"},
	{"x.([]string)", "x.(...)"},
}

func TestExprs(t *testing.T) {
	for _, test := range testExprs {
		src := "package p; var _ = " + test.src + "; var (x, y int; s []string; f func(int, float32))"
		pkg, err := makePkg(t, src)
		if err != nil {
			t.Errorf("%s: %s", src, err)
			continue
		}
		// TODO(gri) writing the code below w/o the decl variable will
		//           cause a 386 compiler error (out of fixed registers)
		decl := pkg.Files[filename].Decls[0].(*ast.GenDecl)
		expr := decl.Specs[0].(*ast.ValueSpec).Values[0]
		str := exprString(expr)
		if str != test.str {
			t.Errorf("%s: got %s, want %s", test.src, str, test.str)
		}
	}
}
<|endoftext|>


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

// Package pkcs12 implements some of PKCS#12.
//
// This implementation is distilled from https://tools.ietf.org/html/rfc7292
// and referenced documents. It is intended for decoding P12/PFX-stored
// certificates and keys for use with the crypto/tls package.
package pkcs12

import (
	"crypto/ecdsa"
	"crypto/rsa"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/asn1"
	"encoding/hex"
	"encoding/pem"
	"errors"
)

var (
	oidDataContentType          = asn1.ObjectIdentifier([]int{1, 2, 840, 113549, 1, 7, 1})
	oidEncryptedDataContentType = asn1.ObjectIdentifier([]int{1, 2, 840, 113549, 1, 7, 6})

	oidFriendlyName     = asn1.ObjectIdentifier([]int{1, 2, 840, 113549, 1, 9, 20})
	oidLocalKeyID       = asn1.ObjectIdentifier([]int{1, 2, 840, 113549, 1, 9, 21})
	oidMicrosoftCSPName = asn1.ObjectIdentifier([]int{1, 3, 6, 1, 4, 1, 311, 17, 1})
)

type pfxPdu struct {
	Version  int
	AuthSafe contentInfo
	MacData  macData `asn1:"optional"`
}

type contentInfo struct {
	ContentType asn1.ObjectIdentifier
	Content     asn1.RawValue `asn1:"tag:0,explicit,optional"`
}

type encryptedData struct {
	Version              int
	EncryptedContentInfo encryptedContentInfo
}

type encryptedContentInfo struct {
	ContentType                asn1.ObjectIdentifier
	ContentEncryptionAlgorithm pkix.AlgorithmIdentifier
	EncryptedContent           []byte `asn1:"tag:0,optional"`
}

func (i encryptedContentInfo) Algorithm() pkix.AlgorithmIdentifier {
	return i.ContentEncryptionAlgorithm
}

func (i encryptedContentInfo) Data() []byte { return i.EncryptedContent }

type safeBag struct {
	Id         asn1.ObjectIdentifier
	Value      asn1.RawValue     `asn1:"tag:0,explicit"`
	Attributes []pkcs12Attribute `asn1:"set,optional"`
}

type pkcs12Attribute struct {
	Id    asn1.ObjectIdentifier
	Value asn1.RawValue `asn1:"set"`
}

type encryptedPrivateKeyInfo struct {
	AlgorithmIdentifier pkix.AlgorithmIdentifier
	EncryptedData       []byte
}

func (i encryptedPrivateKeyInfo) Algorithm() pkix.AlgorithmIdentifier {
	return i.AlgorithmIdentifier
}

func (i encryptedPrivateKeyInfo) Data() []byte {
	return i.EncryptedData
}

// PEM block types
const (
	certificateType = "CERTIFICATE"
	privateKeyType  = "PRIVATE KEY"
)

// unmarshal calls asn1.Unmarshal, but also returns an error if there is any
// trailing data after unmarshaling.
func unmarshal(in []byte, out interface{}) error {
	trailing, err := asn1.Unmarshal(in, out)
	if err != nil {
		return err
	}
	if len(trailing) != 0 {
		return errors.New("pkcs12: trailing data found")
	}
	return nil
}

// ConvertToPEM converts all "safe bags" contained in pfxData to PEM blocks.
func ToPEM(pfxData []byte, password string) ([]*pem.Block, error) {
	encodedPassword, err := bmpString(password)
	if err != nil {
		return nil, ErrIncorrectPassword
	}

	bags, encodedPassword, err := getSafeContents(pfxData, encodedPassword)

	if err != nil {
		return nil, err
	}

	blocks := make([]*pem.Block, 0, len(bags))
	for _, bag := range bags {
		block, err := convertBag(&bag, encodedPassword)
		if err != nil {
			return nil, err
		}
		blocks = append(blocks, block)
	}

	return blocks, nil
}

func convertBag(bag *safeBag, password []byte) (*pem.Block, error) {
	block := &pem.Block{
		Headers: make(map[string]string),
	}

	for _, attribute := range bag.Attributes {
		k, v, err := convertAttribute(&attribute)
		if err != nil {
			return nil, err
		}
		block.Headers[k] = v
	}

	switch {
	case bag.Id.Equal(oidCertBag):
		block.Type = certificateType
		certsData, err := decodeCertBag(bag.Value.Bytes)
		if err != nil {
			return nil, err
		}
		block.Bytes = certsData
	case bag.Id.Equal(oidPKCS8ShroundedKeyBag):
		block.Type = privateKeyType

		key, err := decodePkcs8ShroudedKeyBag(bag.Value.Bytes, password)
		if err != nil {
			return nil, err
		}

		switch key := key.(type) {
		case *rsa.PrivateKey:
			block.Bytes = x509.MarshalPKCS1PrivateKey(key)
		case *ecdsa.PrivateKey:
			block.Bytes, err = x509.MarshalECPrivateKey(key)
			if err != nil {
				return nil, err
			}
		default:
			return nil, errors.New("found unknown private key type in PKCS#8 wrapping")
		}
	default:
		return nil, errors.New("don't know how to convert a safe bag of type " + bag.Id.String())
	}
	return block, nil
}

func convertAttribute(attribute *pkcs12Attribute) (key, value string, err error) {
	isString := false

	switch {
	case attribute.Id.Equal(oidFriendlyName):
		key = "friendlyName"
		isString = true
	case attribute.Id.Equal(oidLocalKeyID):
		key = "localKeyId"
	case attribute.Id.Equal(oidMicrosoftCSPName):
		// This key is chosen to match OpenSSL.
		key = "Microsoft CSP Name"
		isString = true
	default:
		return "", "", errors.New("pkcs12: unknown attribute with OID " + attribute.Id.String())
	}

	if isString {
		if err := unmarshal(attribute.Value.Bytes, &attribute.Value); err != nil {
			return "", "", err
		}
		if value, err = decodeBMPString(attribute.Value.Bytes); err != nil {
			return "", "", err
		}
	} else {
		var id []byte
		if err := unmarshal(attribute.Value.Bytes, &id); err != nil {
			return "", "", err
		}
		value = hex.EncodeToString(id)
	}

	return key, value, nil
}

// Decode extracts a certificate and private key from pfxData. This function
// assumes that there is only one certificate and only one private key in the
// pfxData.
func Decode(pfxData []byte, password string) (privateKey interface{}, certificate *x509.Certificate, err error) {
	encodedPassword, err := bmpString(password)
	if err != nil {
		return nil, nil, err
	}

	bags, encodedPassword, err := getSafeContents(pfxData, encodedPassword)
	if err != nil {
		return nil, nil, err
	}

	if len(bags) != 2 {
		err = errors.New("pkcs12: expected exactly two safe bags in the PFX PDU")
		return
	}

	for _, bag := range bags {
		switch {
		case bag.Id.Equal(oidCertBag):
			if certificate != nil {
				err = errors.New("pkcs12: expected exactly one certificate bag")
			}

			certsData, err := decodeCertBag(bag.Value.Bytes)
			if err != nil {
				return nil, nil, err
			}
			certs, err := x509.ParseCertificates(certsData)
			if err != nil {
				return nil, nil, err
			}
			if len(certs) != 1 {
				err = errors.New("pkcs12: expected exactly one certificate in the certBag")
				return nil, nil, err
			}
			certificate = certs[0]

		case bag.Id.Equal(oidPKCS8ShroundedKeyBag):
			if privateKey != nil {
				err = errors.New("pkcs12: expected exactly one key bag")
			}

			if privateKey, err = decodePkcs8ShroudedKeyBag(bag.Value.Bytes, encodedPassword); err != nil {
				return nil, nil, err
			}
		}
	}

	if certificate == nil {
		return nil, nil, errors.New("pkcs12: certificate missing")
	}
	if privateKey == nil {
		return nil, nil, errors.New("pkcs12: private key missing")
	}

	return
}

// DecodeAllCerts extracts a certificate and private key from pfxData. This function
// assumes that there is only one certificate and only one private key in the
// pfxData.
func DecodeAllCerts(pfxData []byte, password string) (privateKey interface{}, certificates []*x509.Certificate, err error) {
	encodedPassword, err := bmpString(password)
	if err != nil {
		return nil, nil, err
	}

	bags, encodedPassword, err := getSafeContents(pfxData, encodedPassword)
	if err != nil {
		return nil, nil, err
	}

	if len(bags) != 2 {
		err = errors.New("pkcs12: expected exactly two safe bags in the PFX PDU")
		return
	}

	certificates = []*x509.Certificate{}
	for _, bag := range bags {
		switch {
		case bag.Id.Equal(oidCertBag):
			certsData, err := decodeCertBag(bag.Value.Bytes)
			if err != nil {
				return nil, nil, err
			}
			certs, err := x509.ParseCertificates(certsData)
			if err != nil {
				return nil, nil, err
			}
			if len(certs) != 1 {
				err = errors.New("pkcs12: expected exactly one certificate in the certBag")
				return nil, nil, err
			}
			certificates = append(certificates, certs[0])
		case bag.Id.Equal(oidPKCS8ShroundedKeyBag):
			if privateKey != nil {
				err = errors.New("pkcs12: expected exactly one key bag")
			}

			if privateKey, err = decodePkcs8ShroudedKeyBag(bag.Value.Bytes, encodedPassword); err != nil {
				return nil, nil, err
			}
		}
	}

	if certificates == nil || len(certificates) == 0 {
		return nil, nil, errors.New("pkcs12: certificate missing")
	}
	if privateKey == nil {
		return nil, nil, errors.New("pkcs12: private key missing")
	}

	return
}

func getSafeContents(p12Data, password []byte) (bags []safeBag, updatedPassword []byte, err error) {
	pfx := new(pfxPdu)
	if err := unmarshal(p12Data, pfx); err != nil {
		return nil, nil, errors.New("pkcs12: error reading P12 data: " + err.Error())
	}

	if pfx.Version != 3 {
		return nil, nil, NotImplementedError("can only decode v3 PFX PDU's")
	}

	if !pfx.AuthSafe.ContentType.Equal(oidDataContentType) {
		return nil, nil, NotImplementedError("only password-protected PFX is implemented")
	}

	// unmarshal the explicit bytes in the content for type 'data'
	if err := unmarshal(pfx.AuthSafe.Content.Bytes, &pfx.AuthSafe.Content); err != nil {
		return nil, nil, err
	}

	if len(pfx.MacData.Mac.Algorithm.Algorithm) == 0 {
		return nil, nil, errors.New("pkcs12: no MAC in data")
	}

	if err := verifyMac(&pfx.MacData, pfx.AuthSafe.Content.Bytes, password); err != nil {
		if err == ErrIncorrectPassword && len(password) == 2 && password[0] == 0 && password[1] == 0 {
			// some implementations use an empty byte array
			// for the empty string password try one more
			// time with empty-empty password
			password = nil
			err = verifyMac(&pfx.MacData, pfx.AuthSafe.Content.Bytes, password)
		}
		if err != nil {
			return nil, nil, err
		}
	}

	var authenticatedSafe []contentInfo
	if err := unmarshal(pfx.AuthSafe.Content.Bytes, &authenticatedSafe); err != nil {
		return nil, nil, err
	}

	if len(authenticatedSafe) != 2 {
		return nil, nil, NotImplementedError("expected exactly two items in the authenticated safe")
	}

	for _, ci := range authenticatedSafe {
		var data []byte

		switch {
		case ci.ContentType.Equal(oidDataContentType):
			if err := unmarshal(ci.Content.Bytes, &data); err != nil {
				return nil, nil, err
			}
		case ci.ContentType.Equal(oidEncryptedDataContentType):
			var encryptedData encryptedData
			if err := unmarshal(ci.Content.Bytes, &encryptedData); err != nil {
				return nil, nil, err
			}
			if encryptedData.Version != 0 {
				return nil, nil, NotImplementedError("only version 0 of EncryptedData is supported")
			}
			if data, err = pbDecrypt(encryptedData.EncryptedContentInfo, password); err != nil {
				return nil, nil, err
			}
		default:
			return nil, nil, NotImplementedError("only data and encryptedData content types are supported in authenticated safe")
		}

		var safeContents []safeBag
		if err := unmarshal(data, &safeContents); err != nil {
			return nil, nil, err
		}
		bags = append(bags, safeContents...)
	}

	return bags, password, nil
}
<commit_msg>update DecodeAllCerts (#65)<commit_after>// Copyright 2015 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package pkcs12 implements some of PKCS#12.
//
// This implementation is distilled from https://tools.ietf.org/html/rfc7292
// and referenced documents. It is intended for decoding P12/PFX-stored
// certificates and keys for use with the crypto/tls package.
package pkcs12

import (
	"crypto/ecdsa"
	"crypto/rsa"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/asn1"
	"encoding/hex"
	"encoding/pem"
	"errors"
)

var (
	oidDataContentType          = asn1.ObjectIdentifier([]int{1, 2, 840, 113549, 1, 7, 1})
	oidEncryptedDataContentType = asn1.ObjectIdentifier([]int{1, 2, 840, 113549, 1, 7, 6})

	oidFriendlyName     = asn1.ObjectIdentifier([]int{1, 2, 840, 113549, 1, 9, 20})
	oidLocalKeyID       = asn1.ObjectIdentifier([]int{1, 2, 840, 113549, 1, 9, 21})
	oidMicrosoftCSPName = asn1.ObjectIdentifier([]int{1, 3, 6, 1, 4, 1, 311, 17, 1})
)

type pfxPdu struct {
	Version  int
	AuthSafe contentInfo
	MacData  macData `asn1:"optional"`
}

type contentInfo struct {
	ContentType asn1.ObjectIdentifier
	Content     asn1.RawValue `asn1:"tag:0,explicit,optional"`
}

type encryptedData struct {
	Version              int
	EncryptedContentInfo encryptedContentInfo
}

type encryptedContentInfo struct {
	ContentType                asn1.ObjectIdentifier
	ContentEncryptionAlgorithm pkix.AlgorithmIdentifier
	EncryptedContent           []byte `asn1:"tag:0,optional"`
}

func (i encryptedContentInfo) Algorithm() pkix.AlgorithmIdentifier {
	return i.ContentEncryptionAlgorithm
}

func (i encryptedContentInfo) Data() []byte { return i.EncryptedContent }

type safeBag struct {
	Id         asn1.ObjectIdentifier
	Value      asn1.RawValue     `asn1:"tag:0,explicit"`
	Attributes []pkcs12Attribute `asn1:"set,optional"`
}

type pkcs12Attribute struct {
	Id    asn1.ObjectIdentifier
	Value asn1.RawValue `asn1:"set"`
}

type encryptedPrivateKeyInfo struct {
	AlgorithmIdentifier pkix.AlgorithmIdentifier
	EncryptedData       []byte
}

func (i encryptedPrivateKeyInfo) Algorithm() pkix.AlgorithmIdentifier {
	return i.AlgorithmIdentifier
}

func (i encryptedPrivateKeyInfo) Data() []byte {
	return i.EncryptedData
}

// PEM block types
const (
	certificateType = "CERTIFICATE"
	privateKeyType  = "PRIVATE KEY"
)

// unmarshal calls asn1.Unmarshal, but also returns an error if there is any
// trailing data after unmarshaling.
func unmarshal(in []byte, out interface{}) error {
	trailing, err := asn1.Unmarshal(in, out)
	if err != nil {
		return err
	}
	if len(trailing) != 0 {
		return errors.New("pkcs12: trailing data found")
	}
	return nil
}

// ConvertToPEM converts all "safe bags" contained in pfxData to PEM blocks.
func ToPEM(pfxData []byte, password string) ([]*pem.Block, error) {
	encodedPassword, err := bmpString(password)
	if err != nil {
		return nil, ErrIncorrectPassword
	}

	bags, encodedPassword, err := getSafeContents(pfxData, encodedPassword)

	if err != nil {
		return nil, err
	}

	blocks := make([]*pem.Block, 0, len(bags))
	for _, bag := range bags {
		block, err := convertBag(&bag, encodedPassword)
		if err != nil {
			return nil, err
		}
		blocks = append(blocks, block)
	}

	return blocks, nil
}

func convertBag(bag *safeBag, password []byte) (*pem.Block, error) {
	block := &pem.Block{
		Headers: make(map[string]string),
	}

	for _, attribute := range bag.Attributes {
		k, v, err := convertAttribute(&attribute)
		if err != nil {
			return nil, err
		}
		block.Headers[k] = v
	}

	switch {
	case bag.Id.Equal(oidCertBag):
		block.Type = certificateType
		certsData, err := decodeCertBag(bag.Value.Bytes)
		if err != nil {
			return nil, err
		}
		block.Bytes = certsData
	case bag.Id.Equal(oidPKCS8ShroundedKeyBag):
		block.Type = privateKeyType

		key, err := decodePkcs8ShroudedKeyBag(bag.Value.Bytes, password)
		if err != nil {
			return nil, err
		}

		switch key := key.(type) {
		case *rsa.PrivateKey:
			block.Bytes = x509.MarshalPKCS1PrivateKey(key)
		case *ecdsa.PrivateKey:
			block.Bytes, err = x509.MarshalECPrivateKey(key)
			if err != nil {
				return nil, err
			}
		default:
			return nil, errors.New("found unknown private key type in PKCS#8 wrapping")
		}
	default:
		return nil, errors.New("don't know how to convert a safe bag of type " + bag.Id.String())
	}
	return block, nil
}

func convertAttribute(attribute *pkcs12Attribute) (key, value string, err error) {
	isString := false

	switch {
	case attribute.Id.Equal(oidFriendlyName):
		key = "friendlyName"
		isString = true
	case attribute.Id.Equal(oidLocalKeyID):
		key = "localKeyId"
	case attribute.Id.Equal(oidMicrosoftCSPName):
		// This key is chosen to match OpenSSL.
		key = "Microsoft CSP Name"
		isString = true
	default:
		return "", "", errors.New("pkcs12: unknown attribute with OID " + attribute.Id.String())
	}

	if isString {
		if err := unmarshal(attribute.Value.Bytes, &attribute.Value); err != nil {
			return "", "", err
		}
		if value, err = decodeBMPString(attribute.Value.Bytes); err != nil {
			return "", "", err
		}
	} else {
		var id []byte
		if err := unmarshal(attribute.Value.Bytes, &id); err != nil {
			return "", "", err
		}
		value = hex.EncodeToString(id)
	}

	return key, value, nil
}

// Decode extracts a certificate and private key from pfxData. This function
// assumes that there is only one certificate and only one private key in the
// pfxData.
func Decode(pfxData []byte, password string) (privateKey interface{}, certificate *x509.Certificate, err error) {
	encodedPassword, err := bmpString(password)
	if err != nil {
		return nil, nil, err
	}

	bags, encodedPassword, err := getSafeContents(pfxData, encodedPassword)
	if err != nil {
		return nil, nil, err
	}

	if len(bags) != 2 {
		err = errors.New("pkcs12: expected exactly two safe bags in the PFX PDU")
		return
	}

	for _, bag := range bags {
		switch {
		case bag.Id.Equal(oidCertBag):
			if certificate != nil {
				err = errors.New("pkcs12: expected exactly one certificate bag")
			}

			certsData, err := decodeCertBag(bag.Value.Bytes)
			if err != nil {
				return nil, nil, err
			}
			certs, err := x509.ParseCertificates(certsData)
			if err != nil {
				return nil, nil, err
			}
			if len(certs) != 1 {
				err = errors.New("pkcs12: expected exactly one certificate in the certBag")
				return nil, nil, err
			}
			certificate = certs[0]

		case bag.Id.Equal(oidPKCS8ShroundedKeyBag):
			if privateKey != nil {
				err = errors.New("pkcs12: expected exactly one key bag")
			}

			if privateKey, err = decodePkcs8ShroudedKeyBag(bag.Value.Bytes, encodedPassword); err != nil {
				return nil, nil, err
			}
		}
	}

	if certificate == nil {
		return nil, nil, errors.New("pkcs12: certificate missing")
	}
	if privateKey == nil {
		return nil, nil, errors.New("pkcs12: private key missing")
	}

	return
}

// DecodeAllCerts extracts a certificates from pfxData.
func DecodeAllCerts(pfxData []byte, password string) (certificates []*x509.Certificate, err error) {
	encodedPassword, err := bmpString(password)
	if err != nil {
		return nil, err
	}

	bags, encodedPassword, err := getSafeContents(pfxData, encodedPassword)
	if err != nil {
		return nil, err
	}

	certificates = []*x509.Certificate{}
	for _, bag := range bags {
		switch {
		case bag.Id.Equal(oidCertBag):
			certsData, err := decodeCertBag(bag.Value.Bytes)
			if err != nil {
				return nil, err
			}
			certs, err := x509.ParseCertificates(certsData)
			if err != nil {
				return nil, err
			}
			if len(certs) != 1 {
				err = errors.New("pkcs12: expected exactly one certificate in the certBag")
				return nil, err
			}
			certificates = append(certificates, certs[0])
		}
	}

	if certificates == nil || len(certificates) == 0 {
		return nil, errors.New("pkcs12: certificate missing")
	}

	return
}

func getSafeContents(p12Data, password []byte) (bags []safeBag, updatedPassword []byte, err error) {
	pfx := new(pfxPdu)
	if err := unmarshal(p12Data, pfx); err != nil {
		return nil, nil, errors.New("pkcs12: error reading P12 data: " + err.Error())
	}

	if pfx.Version != 3 {
		return nil, nil, NotImplementedError("can only decode v3 PFX PDU's")
	}

	if !pfx.AuthSafe.ContentType.Equal(oidDataContentType) {
		return nil, nil, NotImplementedError("only password-protected PFX is implemented")
	}

	// unmarshal the explicit bytes in the content for type 'data'
	if err := unmarshal(pfx.AuthSafe.Content.Bytes, &pfx.AuthSafe.Content); err != nil {
		return nil, nil, err
	}

	if len(pfx.MacData.Mac.Algorithm.Algorithm) == 0 {
		return nil, nil, errors.New("pkcs12: no MAC in data")
	}

	if err := verifyMac(&pfx.MacData, pfx.AuthSafe.Content.Bytes, password); err != nil {
		if err == ErrIncorrectPassword && len(password) == 2 && password[0] == 0 && password[1] == 0 {
			// some implementations use an empty byte array
			// for the empty string password try one more
			// time with empty-empty password
			password = nil
			err = verifyMac(&pfx.MacData, pfx.AuthSafe.Content.Bytes, password)
		}
		if err != nil {
			return nil, nil, err
		}
	}

	var authenticatedSafe []contentInfo
	if err := unmarshal(pfx.AuthSafe.Content.Bytes, &authenticatedSafe); err != nil {
		return nil, nil, err
	}

	if len(authenticatedSafe) != 2 {
		return nil, nil, NotImplementedError("expected exactly two items in the authenticated safe")
	}

	for _, ci := range authenticatedSafe {
		var data []byte

		switch {
		case ci.ContentType.Equal(oidDataContentType):
			if err := unmarshal(ci.Content.Bytes, &data); err != nil {
				return nil, nil, err
			}
		case ci.ContentType.Equal(oidEncryptedDataContentType):
			var encryptedData encryptedData
			if err := unmarshal(ci.Content.Bytes, &encryptedData); err != nil {
				return nil, nil, err
			}
			if encryptedData.Version != 0 {
				return nil, nil, NotImplementedError("only version 0 of EncryptedData is supported")
			}
			if data, err = pbDecrypt(encryptedData.EncryptedContentInfo, password); err != nil {
				return nil, nil, err
			}
		default:
			return nil, nil, NotImplementedError("only data and encryptedData content types are supported in authenticated safe")
		}

		var safeContents []safeBag
		if err := unmarshal(data, &safeContents); err != nil {
			return nil, nil, err
		}
		bags = append(bags, safeContents...)
	}

	return bags, password, nil
}
<|endoftext|>


<commit_before>// Package acme provides a means of performing Let's Encrypt DNS challenges via a DNSConfig
package acme

import (
	"crypto/x509"
	"encoding/pem"
	"fmt"
	"io/ioutil"
	"log"
	"net/url"
	"sort"
	"strings"
	"time"

	"github.com/StackExchange/dnscontrol/models"
	"github.com/StackExchange/dnscontrol/pkg/nameservers"
	"github.com/StackExchange/dnscontrol/pkg/notifications"
	"github.com/go-acme/lego/certcrypto"
	"github.com/go-acme/lego/certificate"
	"github.com/go-acme/lego/challenge"
	"github.com/go-acme/lego/challenge/dns01"
	"github.com/go-acme/lego/lego"
	acmelog "github.com/go-acme/lego/log"
)

type CertConfig struct {
	CertName   string   `json:"cert_name"`
	Names      []string `json:"names"`
	UseECC     bool     `json:"use_ecc"`
	MustStaple bool     `json:"must_staple"`
}

type Client interface {
	IssueOrRenewCert(config *CertConfig, renewUnder int, verbose bool) (bool, error)
}

type certManager struct {
	email         string
	acmeDirectory string
	acmeHost      string

	storage         Storage
	cfg             *models.DNSConfig
	domains         map[string]*models.DomainConfig
	originalDomains []*models.DomainConfig

	notifier notifications.Notifier

	account    *Account
	waitedOnce bool
}

const (
	LetsEncryptLive  = "https://acme-v02.api.letsencrypt.org/directory"
	LetsEncryptStage = "https://acme-staging-v02.api.letsencrypt.org/directory"
)

func New(cfg *models.DNSConfig, directory string, email string, server string, notify notifications.Notifier) (Client, error) {
	return commonNew(cfg, directoryStorage(directory), email, server, notify)
}

func commonNew(cfg *models.DNSConfig, storage Storage, email string, server string, notify notifications.Notifier) (Client, error) {
	u, err := url.Parse(server)
	if err != nil || u.Host == "" {
		return nil, fmt.Errorf("ACME directory '%s' is not a valid URL", server)
	}
	c := &certManager{
		storage:       storage,
		email:         email,
		acmeDirectory: server,
		acmeHost:      u.Host,
		cfg:           cfg,
		domains:       map[string]*models.DomainConfig{},
		notifier:      notify,
	}

	acct, err := c.getOrCreateAccount()
	if err != nil {
		return nil, err
	}
	c.account = acct
	return c, nil
}

func NewVault(cfg *models.DNSConfig, vaultPath string, email string, server string, notify notifications.Notifier) (Client, error) {
	storage, err := makeVaultStorage(vaultPath)
	if err != nil {
		return nil, err
	}
	return commonNew(cfg, storage, email, server, notify)
}

// IssueOrRenewCert will obtain a certificate with the given name if it does not exist,
// or renew it if it is close enough to the expiration date.
// It will return true if it issued or updated the certificate.
func (c *certManager) IssueOrRenewCert(cfg *CertConfig, renewUnder int, verbose bool) (bool, error) {
	if !verbose {
		acmelog.Logger = log.New(ioutil.Discard, "", 0)
	}
	defer c.finalCleanUp()

	log.Printf("Checking certificate [%s]", cfg.CertName)
	existing, err := c.storage.GetCertificate(cfg.CertName)
	if err != nil {
		return false, err
	}

	var client *lego.Client

	var action = func() (*certificate.Resource, error) {
		return client.Certificate.Obtain(certificate.ObtainRequest{
			Bundle:     true,
			Domains:    cfg.Names,
			MustStaple: cfg.MustStaple,
		})
	}

	if existing == nil {
		log.Println("No existing cert found. Issuing new...")
	} else {
		names, daysLeft, err := getCertInfo(existing.Certificate)
		if err != nil {
			return false, err
		}
		log.Printf("Found existing cert. %0.2f days remaining.", daysLeft)
		namesOK := dnsNamesEqual(cfg.Names, names)
		if daysLeft >= float64(renewUnder) && namesOK {
			log.Println("Nothing to do")
			//nothing to do
			return false, nil
		}
		if !namesOK {
			log.Println("DNS Names don't match expected set. Reissuing.")
		} else {
			log.Println("Renewing cert")
			action = func() (*certificate.Resource, error) {
				return client.Certificate.Renew(*existing, true, cfg.MustStaple)
			}
		}
	}

	kt := certcrypto.RSA2048
	if cfg.UseECC {
		kt = certcrypto.EC256
	}
	config := lego.NewConfig(c.account)
	config.CADirURL = c.acmeDirectory
	config.Certificate.KeyType = kt
	client, err = lego.NewClient(config)
	if err != nil {
		return false, err
	}
	client.Challenge.Remove(challenge.HTTP01)
	client.Challenge.Remove(challenge.TLSALPN01)
	client.Challenge.SetDNS01Provider(c)

	dns01.WrapPreCheck(c.preCheckDNS)

	certResource, err := action()
	if err != nil {
		return false, err
	}
	fmt.Printf("Obtained certificate for %s\n", cfg.CertName)
	if err = c.storage.StoreCertificate(cfg.CertName, certResource); err != nil {
		return true, err
	}

	return true, nil
}

func getCertInfo(pemBytes []byte) (names []string, remaining float64, err error) {
	block, _ := pem.Decode(pemBytes)
	if block == nil {
		return nil, 0, fmt.Errorf("Invalid certificate pem data")
	}
	cert, err := x509.ParseCertificate(block.Bytes)
	if err != nil {
		return nil, 0, err
	}
	var daysLeft = float64(cert.NotAfter.Sub(time.Now())) / float64(time.Hour*24)
	return cert.DNSNames, daysLeft, nil
}

// checks two lists of sans to make sure they have all the same names in them.
func dnsNamesEqual(a []string, b []string) bool {
	if len(a) != len(b) {
		return false
	}
	sort.Strings(a)
	sort.Strings(b)
	for i, s := range a {
		if b[i] != s {
			return false
		}
	}
	return true
}

func (c *certManager) Present(domain, token, keyAuth string) (e error) {
	d := c.cfg.DomainContainingFQDN(domain)
	name := d.Name
	if seen := c.domains[name]; seen != nil {
		// we've already pre-processed this domain, just need to add to it.
		d = seen
	} else {
		// one-time tasks to get this domain ready.
		// if multiple validations on a single domain, we don't need to rebuild all this.

		// fix NS records for this domain's DNS providers
		nsList, err := nameservers.DetermineNameservers(d)
		if err != nil {
			return err
		}
		d.Nameservers = nsList
		nameservers.AddNSRecords(d)

		// make sure we have the latest config before we change anything.
		// alternately, we could avoid a lot of this trouble if we really really trusted no-purge in all cases
		if err := c.ensureNoPendingCorrections(d); err != nil {
			return err
		}

		// copy domain and work from copy from now on. That way original config can be used to "restore" when we are all done.
		copy, err := d.Copy()
		if err != nil {
			return err
		}
		c.originalDomains = append(c.originalDomains, d)
		c.domains[name] = copy
		d = copy
	}

	fqdn, val := dns01.GetRecord(domain, keyAuth)
	txt := &models.RecordConfig{Type: "TXT"}
	txt.SetTargetTXT(val)
	txt.SetLabelFromFQDN(fqdn, d.Name)
	d.Records = append(d.Records, txt)
	return c.getAndRunCorrections(d)
}

func (c *certManager) ensureNoPendingCorrections(d *models.DomainConfig) error {
	corrections, err := c.getCorrections(d)
	if err != nil {
		return err
	}
	if len(corrections) != 0 {
		// TODO: maybe allow forcing through this check.
		for _, c := range corrections {
			fmt.Println(c.Msg)
		}
		return fmt.Errorf("Found %d pending corrections for %s. Not going to proceed issuing certificates", len(corrections), d.Name)
	}
	return nil
}

// IgnoredProviders is a lit of provider names that should not be used to fill challenges.
var IgnoredProviders = map[string]bool{}

func (c *certManager) getCorrections(d *models.DomainConfig) ([]*models.Correction, error) {
	cs := []*models.Correction{}
	for _, p := range d.DNSProviderInstances {
		if IgnoredProviders[p.Name] {
			continue
		}
		dc, err := d.Copy()
		if err != nil {
			return nil, err
		}
		corrections, err := p.Driver.GetDomainCorrections(dc)
		if err != nil {
			return nil, err
		}
		for _, c := range corrections {
			c.Msg = fmt.Sprintf("[%s] %s", p.Name, strings.TrimSpace(c.Msg))
		}
		cs = append(cs, corrections...)
	}
	return cs, nil
}

func (c *certManager) getAndRunCorrections(d *models.DomainConfig) error {
	cs, err := c.getCorrections(d)
	if err != nil {
		return err
	}
	fmt.Printf("%d corrections\n", len(cs))
	for _, corr := range cs {
		fmt.Printf("Running [%s]\n", corr.Msg)
		err = corr.F()
		c.notifier.Notify(d.Name, "certs", corr.Msg, err, false)
		if err != nil {
			return err
		}
	}
	return nil
}

func (c *certManager) CleanUp(domain, token, keyAuth string) error {
	// do nothing for now. We will do a final clean up step at the very end.
	return nil
}

func (c *certManager) finalCleanUp() error {
	log.Println("Cleaning up all records we made")
	var lastError error
	for _, d := range c.originalDomains {
		if err := c.getAndRunCorrections(d); err != nil {
			log.Printf("ERROR cleaning up: %s", err)
			lastError = err
		}
	}
	return lastError
}
<commit_msg>get-certs: DNS01 challenge skipping preCheckDNS (#591)<commit_after>// Package acme provides a means of performing Let's Encrypt DNS challenges via a DNSConfig
package acme

import (
	"crypto/x509"
	"encoding/pem"
	"fmt"
	"io/ioutil"
	"log"
	"net/url"
	"sort"
	"strings"
	"time"

	"github.com/StackExchange/dnscontrol/models"
	"github.com/StackExchange/dnscontrol/pkg/nameservers"
	"github.com/StackExchange/dnscontrol/pkg/notifications"
	"github.com/go-acme/lego/certcrypto"
	"github.com/go-acme/lego/certificate"
	"github.com/go-acme/lego/challenge"
	"github.com/go-acme/lego/challenge/dns01"
	"github.com/go-acme/lego/lego"
	acmelog "github.com/go-acme/lego/log"
)

type CertConfig struct {
	CertName   string   `json:"cert_name"`
	Names      []string `json:"names"`
	UseECC     bool     `json:"use_ecc"`
	MustStaple bool     `json:"must_staple"`
}

type Client interface {
	IssueOrRenewCert(config *CertConfig, renewUnder int, verbose bool) (bool, error)
}

type certManager struct {
	email         string
	acmeDirectory string
	acmeHost      string

	storage         Storage
	cfg             *models.DNSConfig
	domains         map[string]*models.DomainConfig
	originalDomains []*models.DomainConfig

	notifier notifications.Notifier

	account    *Account
	waitedOnce bool
}

const (
	LetsEncryptLive  = "https://acme-v02.api.letsencrypt.org/directory"
	LetsEncryptStage = "https://acme-staging-v02.api.letsencrypt.org/directory"
)

func New(cfg *models.DNSConfig, directory string, email string, server string, notify notifications.Notifier) (Client, error) {
	return commonNew(cfg, directoryStorage(directory), email, server, notify)
}

func commonNew(cfg *models.DNSConfig, storage Storage, email string, server string, notify notifications.Notifier) (Client, error) {
	u, err := url.Parse(server)
	if err != nil || u.Host == "" {
		return nil, fmt.Errorf("ACME directory '%s' is not a valid URL", server)
	}
	c := &certManager{
		storage:       storage,
		email:         email,
		acmeDirectory: server,
		acmeHost:      u.Host,
		cfg:           cfg,
		domains:       map[string]*models.DomainConfig{},
		notifier:      notify,
	}

	acct, err := c.getOrCreateAccount()
	if err != nil {
		return nil, err
	}
	c.account = acct
	return c, nil
}

func NewVault(cfg *models.DNSConfig, vaultPath string, email string, server string, notify notifications.Notifier) (Client, error) {
	storage, err := makeVaultStorage(vaultPath)
	if err != nil {
		return nil, err
	}
	return commonNew(cfg, storage, email, server, notify)
}

// IssueOrRenewCert will obtain a certificate with the given name if it does not exist,
// or renew it if it is close enough to the expiration date.
// It will return true if it issued or updated the certificate.
func (c *certManager) IssueOrRenewCert(cfg *CertConfig, renewUnder int, verbose bool) (bool, error) {
	if !verbose {
		acmelog.Logger = log.New(ioutil.Discard, "", 0)
	}
	defer c.finalCleanUp()

	log.Printf("Checking certificate [%s]", cfg.CertName)
	existing, err := c.storage.GetCertificate(cfg.CertName)
	if err != nil {
		return false, err
	}

	var client *lego.Client

	var action = func() (*certificate.Resource, error) {
		return client.Certificate.Obtain(certificate.ObtainRequest{
			Bundle:     true,
			Domains:    cfg.Names,
			MustStaple: cfg.MustStaple,
		})
	}

	if existing == nil {
		log.Println("No existing cert found. Issuing new...")
	} else {
		names, daysLeft, err := getCertInfo(existing.Certificate)
		if err != nil {
			return false, err
		}
		log.Printf("Found existing cert. %0.2f days remaining.", daysLeft)
		namesOK := dnsNamesEqual(cfg.Names, names)
		if daysLeft >= float64(renewUnder) && namesOK {
			log.Println("Nothing to do")
			//nothing to do
			return false, nil
		}
		if !namesOK {
			log.Println("DNS Names don't match expected set. Reissuing.")
		} else {
			log.Println("Renewing cert")
			action = func() (*certificate.Resource, error) {
				return client.Certificate.Renew(*existing, true, cfg.MustStaple)
			}
		}
	}

	kt := certcrypto.RSA2048
	if cfg.UseECC {
		kt = certcrypto.EC256
	}
	config := lego.NewConfig(c.account)
	config.CADirURL = c.acmeDirectory
	config.Certificate.KeyType = kt
	client, err = lego.NewClient(config)
	if err != nil {
		return false, err
	}
	client.Challenge.Remove(challenge.HTTP01)
	client.Challenge.Remove(challenge.TLSALPN01)
	client.Challenge.SetDNS01Provider(c, dns01.WrapPreCheck(c.preCheckDNS))

	certResource, err := action()
	if err != nil {
		return false, err
	}
	fmt.Printf("Obtained certificate for %s\n", cfg.CertName)
	if err = c.storage.StoreCertificate(cfg.CertName, certResource); err != nil {
		return true, err
	}

	return true, nil
}

func getCertInfo(pemBytes []byte) (names []string, remaining float64, err error) {
	block, _ := pem.Decode(pemBytes)
	if block == nil {
		return nil, 0, fmt.Errorf("Invalid certificate pem data")
	}
	cert, err := x509.ParseCertificate(block.Bytes)
	if err != nil {
		return nil, 0, err
	}
	var daysLeft = float64(cert.NotAfter.Sub(time.Now())) / float64(time.Hour*24)
	return cert.DNSNames, daysLeft, nil
}

// checks two lists of sans to make sure they have all the same names in them.
func dnsNamesEqual(a []string, b []string) bool {
	if len(a) != len(b) {
		return false
	}
	sort.Strings(a)
	sort.Strings(b)
	for i, s := range a {
		if b[i] != s {
			return false
		}
	}
	return true
}

func (c *certManager) Present(domain, token, keyAuth string) (e error) {
	d := c.cfg.DomainContainingFQDN(domain)
	name := d.Name
	if seen := c.domains[name]; seen != nil {
		// we've already pre-processed this domain, just need to add to it.
		d = seen
	} else {
		// one-time tasks to get this domain ready.
		// if multiple validations on a single domain, we don't need to rebuild all this.

		// fix NS records for this domain's DNS providers
		nsList, err := nameservers.DetermineNameservers(d)
		if err != nil {
			return err
		}
		d.Nameservers = nsList
		nameservers.AddNSRecords(d)

		// make sure we have the latest config before we change anything.
		// alternately, we could avoid a lot of this trouble if we really really trusted no-purge in all cases
		if err := c.ensureNoPendingCorrections(d); err != nil {
			return err
		}

		// copy domain and work from copy from now on. That way original config can be used to "restore" when we are all done.
		copy, err := d.Copy()
		if err != nil {
			return err
		}
		c.originalDomains = append(c.originalDomains, d)
		c.domains[name] = copy
		d = copy
	}

	fqdn, val := dns01.GetRecord(domain, keyAuth)
	txt := &models.RecordConfig{Type: "TXT"}
	txt.SetTargetTXT(val)
	txt.SetLabelFromFQDN(fqdn, d.Name)
	d.Records = append(d.Records, txt)
	return c.getAndRunCorrections(d)
}

func (c *certManager) ensureNoPendingCorrections(d *models.DomainConfig) error {
	corrections, err := c.getCorrections(d)
	if err != nil {
		return err
	}
	if len(corrections) != 0 {
		// TODO: maybe allow forcing through this check.
		for _, c := range corrections {
			fmt.Println(c.Msg)
		}
		return fmt.Errorf("Found %d pending corrections for %s. Not going to proceed issuing certificates", len(corrections), d.Name)
	}
	return nil
}

// IgnoredProviders is a lit of provider names that should not be used to fill challenges.
var IgnoredProviders = map[string]bool{}

func (c *certManager) getCorrections(d *models.DomainConfig) ([]*models.Correction, error) {
	cs := []*models.Correction{}
	for _, p := range d.DNSProviderInstances {
		if IgnoredProviders[p.Name] {
			continue
		}
		dc, err := d.Copy()
		if err != nil {
			return nil, err
		}
		corrections, err := p.Driver.GetDomainCorrections(dc)
		if err != nil {
			return nil, err
		}
		for _, c := range corrections {
			c.Msg = fmt.Sprintf("[%s] %s", p.Name, strings.TrimSpace(c.Msg))
		}
		cs = append(cs, corrections...)
	}
	return cs, nil
}

func (c *certManager) getAndRunCorrections(d *models.DomainConfig) error {
	cs, err := c.getCorrections(d)
	if err != nil {
		return err
	}
	fmt.Printf("%d corrections\n", len(cs))
	for _, corr := range cs {
		fmt.Printf("Running [%s]\n", corr.Msg)
		err = corr.F()
		c.notifier.Notify(d.Name, "certs", corr.Msg, err, false)
		if err != nil {
			return err
		}
	}
	return nil
}

func (c *certManager) CleanUp(domain, token, keyAuth string) error {
	// do nothing for now. We will do a final clean up step at the very end.
	return nil
}

func (c *certManager) finalCleanUp() error {
	log.Println("Cleaning up all records we made")
	var lastError error
	for _, d := range c.originalDomains {
		if err := c.getAndRunCorrections(d); err != nil {
			log.Printf("ERROR cleaning up: %s", err)
			lastError = err
		}
	}
	return lastError
}
<|endoftext|>


<commit_before>package apl

import "k8s.io/kubernetes/test/e2e_node/services"

// CreateResult is the PK of the row just created
// Or it's a more complex struct of results.
type CreateResult struct {
	Data interface{} `json:"data"`
}

// ModifyResult is the information returned after a delete/update
type ModifyResult struct {
	Skipped   int `json:"skipped"`
	Deleted   int `json:"deleted"`
	Unchanged int `json:"unchanged"`
	Errors    int `json:"errors"`
	Replaced  int `json:"replaced"`
	Inserted  int `json:"inserted"`
} // `json:"data"`

// ModifyOutput wraps a ModifyResult with 'data'
type ModifyOutput struct {
	ModifyResult `json:"data"`
}

// CreateInput Used to wrap the data json required by the api
type CreateInput struct {
	Data interface{} `json:"data"`
}

// UpdateInput Used to wrap the data json required by the api
type UpdateInput struct {
	Data interface{} `json:"data"`
}

// ListParams Used to wrap params for function passing
type ListParams struct {
	Data interface{} `json:"data"`
}

// CreatedByUser used as a join in many results
type CreatedByUser struct {
	ID        string `json:"id,omitempty"`
	FirstName string `json:"first_name,omitempty"`
	LastName  string `json:"last_name,omitempty"`
	Email     string `json:"email,omitempty"`
} // `json:"created_by_user"`

// MetaData ...
type MetaData struct {
	Link        string `json:"link,omitempty"`
	DisplayName string `json:"display_name,omitempty"`
	Description string `json:"description,omitempty"`
	Icon        string `json:"icon,omitempty"`
} // `json:"meta_data"`

// Releases ...
type Releases struct {
	ID       string `json:"id"`
	Version  int    `json:"version"`
	MetaData `json:"meta_data"`
} // `json:"releases"`

// Artifact
type Artifact struct {
	Code    interface{} `json:"code,omitempty"`
	Config  string      `json:"config,omitempty"`
	Image   string      `json:"image,omitempty"`
	Data    string      `json:"data,omitempty"`
	Builder string      `json:"builder,omitempty"`
} // `json:"artifacts"`

// Build
type Build struct {
	Artifact  `json:"artifacts"`
	BuildVars `json:"buildvars,omitempty"`
} // `json:"build,omitempty"`

// Overrides
type Overrides struct {
	Build `json:"build,omitempty"`
} // `json:"overrides,omitempty"`

// Service
type Service struct {
	Name      string      `json:"name,omitempty"`
	Build     interface{} `json:"build,omitempty"`
	Run       interface{} `json:"run,omitempty"`
	Overrides interface{} `json:"overrides,omitempty"`
}

// BuildVars
type BuildVars []map[string]string // `json:"buildvars,omitempty"`

// EnvVars
type EnvVars []map[string]string // `json:"envvars,omitempty"`

// Run
type Run struct {
	Instances int `json:"instances,omitempty"`
} // `json:"run,omitempty"`

//// HealthProbe ...
//type HealthProbe   struct {
//	Delay   interface{} `json:"delay,omitempty"`
//	URL     interface{} `json:"url,omitempty"`
//	Type    interface{} `json:"type,omitempty"`
//	Port    interface{} `json:"port,omitempty"`
//	Timeout interface{} `json:"timeout,omitempty"`
//} // `json:"health_probe"`
//
//// Defaults ...
//type Defaults struct {
//	// TODO: Figure out why this is an int and not string
//	//ReserveMemory   string `json:"reserve_memory,omitempty"`
//	Instances       int    `json:"instances,omitempty"`
//	StorageSize     int    `json:"storage_size,omitempty"`
//	ClientPort      int    `json:"client_port,omitempty"`
//	LimitMemory     string `json:"limit_memory,omitempty"`
//	StorageType     string `json:"storage_type,omitempty"`
//	HealthCheckType string `json:"health_check_type,omitempty"`
//	ServicePort     int    `json:"service_port,omitempty"`
//	ServiceName     string `json:"service_name,omitempty"`
//	ReserveCPU      string `json:"reserve_cpu,omitempty"`
//	ExternalNetwork bool   `json:"external_network,omitempty"`
//	LimitCPU        string `json:"limit_cpu,omitempty"`
//	HealthProbe     `json:"health_probe,omitempty"`
//} // `json:"defaults"`
<commit_msg>rogue import<commit_after>package apl

// CreateResult is the PK of the row just created
// Or it's a more complex struct of results.
type CreateResult struct {
	Data interface{} `json:"data"`
}

// ModifyResult is the information returned after a delete/update
type ModifyResult struct {
	Skipped   int `json:"skipped"`
	Deleted   int `json:"deleted"`
	Unchanged int `json:"unchanged"`
	Errors    int `json:"errors"`
	Replaced  int `json:"replaced"`
	Inserted  int `json:"inserted"`
} // `json:"data"`

// ModifyOutput wraps a ModifyResult with 'data'
type ModifyOutput struct {
	ModifyResult `json:"data"`
}

// CreateInput Used to wrap the data json required by the api
type CreateInput struct {
	Data interface{} `json:"data"`
}

// UpdateInput Used to wrap the data json required by the api
type UpdateInput struct {
	Data interface{} `json:"data"`
}

// ListParams Used to wrap params for function passing
type ListParams struct {
	Data interface{} `json:"data"`
}

// CreatedByUser used as a join in many results
type CreatedByUser struct {
	ID        string `json:"id,omitempty"`
	FirstName string `json:"first_name,omitempty"`
	LastName  string `json:"last_name,omitempty"`
	Email     string `json:"email,omitempty"`
} // `json:"created_by_user"`

// MetaData ...
type MetaData struct {
	Link        string `json:"link,omitempty"`
	DisplayName string `json:"display_name,omitempty"`
	Description string `json:"description,omitempty"`
	Icon        string `json:"icon,omitempty"`
} // `json:"meta_data"`

// Releases ...
type Releases struct {
	ID       string `json:"id"`
	Version  int    `json:"version"`
	MetaData `json:"meta_data"`
} // `json:"releases"`

// Artifact
type Artifact struct {
	Code    interface{} `json:"code,omitempty"`
	Config  string      `json:"config,omitempty"`
	Image   string      `json:"image,omitempty"`
	Data    string      `json:"data,omitempty"`
	Builder string      `json:"builder,omitempty"`
} // `json:"artifacts"`

// Build
type Build struct {
	Artifact  `json:"artifacts"`
	BuildVars `json:"buildvars,omitempty"`
} // `json:"build,omitempty"`

// Overrides
type Overrides struct {
	Build `json:"build,omitempty"`
} // `json:"overrides,omitempty"`

// Service
type Service struct {
	Name      string      `json:"name,omitempty"`
	Build     interface{} `json:"build,omitempty"`
	Run       interface{} `json:"run,omitempty"`
	Overrides interface{} `json:"overrides,omitempty"`
}

// BuildVars
type BuildVars []map[string]string // `json:"buildvars,omitempty"`

// EnvVars
type EnvVars []map[string]string // `json:"envvars,omitempty"`

// Run
type Run struct {
	Instances int `json:"instances,omitempty"`
} // `json:"run,omitempty"`

//// HealthProbe ...
//type HealthProbe   struct {
//	Delay   interface{} `json:"delay,omitempty"`
//	URL     interface{} `json:"url,omitempty"`
//	Type    interface{} `json:"type,omitempty"`
//	Port    interface{} `json:"port,omitempty"`
//	Timeout interface{} `json:"timeout,omitempty"`
//} // `json:"health_probe"`
//
//// Defaults ...
//type Defaults struct {
//	// TODO: Figure out why this is an int and not string
//	//ReserveMemory   string `json:"reserve_memory,omitempty"`
//	Instances       int    `json:"instances,omitempty"`
//	StorageSize     int    `json:"storage_size,omitempty"`
//	ClientPort      int    `json:"client_port,omitempty"`
//	LimitMemory     string `json:"limit_memory,omitempty"`
//	StorageType     string `json:"storage_type,omitempty"`
//	HealthCheckType string `json:"health_check_type,omitempty"`
//	ServicePort     int    `json:"service_port,omitempty"`
//	ServiceName     string `json:"service_name,omitempty"`
//	ReserveCPU      string `json:"reserve_cpu,omitempty"`
//	ExternalNetwork bool   `json:"external_network,omitempty"`
//	LimitCPU        string `json:"limit_cpu,omitempty"`
//	HealthProbe     `json:"health_probe,omitempty"`
//} // `json:"defaults"`
<|endoftext|>


<commit_before>// Package dexec is a logging variant of os/exec.
//
// dexec is *almost* a drop-in replacement for os/exec.  Differences
// are:
//
// - The "Command" function is missing, because a context is always
// required; use CommandContext.
//
// - It is not valid to create a "Cmd" entirely by hand; you must
// create it using CommandContext.  After it has been created, you may
// adjust the fields as you would with an os/exec.Cmd.
//
// The logger used is configured in the context.Context passed to
// CommandContext by calling
// github.com/datawire/ambassador/pkg/dlog.WithLogger.
//
// A Cmd logs when it starts, its exit status, and if they aren't an
// *os.File, logs everything read from or written to .Stdin, .Stdout,
// and .Stderr.  If one of those is an *os.File (as it is following a
// call to .StdinPipe, .StdoutPipe, or .StderrPipe), then that stream
// won't be logged (but it will print a message at process-start
// noting that it isn't being logged).
//
// For example:
//
//     ctx := dlog.WithLogger(context.Background(), myLogger)
//     cmd := dexec.CommandContext(ctx, "printf", "%s\n", "foo bar", "baz")
//     cmd.Stdin = os.Stdin
//     err := cmd.Run()
//
// will log the lines
//
//     [pid:24272] started command []string{"printf", "%s\n", "foo bar", "baz"}
//     [pid:24272] stdin  < not logging input read from file /dev/stdin
//     [pid:24272] stdout+stderr > "foo bar\n"
//     [pid:24272] stdout+stderr > "baz\n"
//     [pid:24272] finished successfully: exit status 0
//
// If you would like a "pipe" to be logged, use an io.Pipe instead of
// calling .StdinPipe, .StdoutPipe, or .StderrPipe.
//
// See the os/exec documentation for more information.
package dexec

import (
	"context"
	"io"
	"os"
	"os/exec"
	"sync"

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

// Error is returned by LookPath when it fails to classify a file as an
// executable.
type Error = exec.Error

// An ExitError reports an unsuccessful exit by a command.
type ExitError = exec.ExitError

// ErrNotFound is the os/exec.ErrNotFound value.
var ErrNotFound = exec.ErrNotFound

// LookPath is the os/exe.LookPath function.
var LookPath = exec.LookPath

// Cmd represents an external command being prepared or run.
//
// A Cmd cannot be reused after calling its Run, Output or CombinedOutput
// methods.
//
// See the os/exec.Cmd documentation for information on the fields
// within it.
//
// Unlike an os/exec.Cmd, you MUST NOT construct a Cmd by hand, it
// must be created with CommandContext.
type Cmd struct {
	*exec.Cmd
	logger dlog.Logger

	pidlock sync.RWMutex
}

// CommandContext returns the Cmd struct to execute the named program with
// the given arguments.
//
// The provided context is used for two purposes:
//
//  1. To kill the process (by calling os.Process.Kill) if the context
//     becomes done before the command completes on its own.
//  2. To get the logger (by calling
//     github.com/datawire/ambassador/pkg/dlog.GetLogger on it).
//
// See the os/exec.Command and os/exec.CommandContext documentation
// for more information.
func CommandContext(ctx context.Context, name string, arg ...string) *Cmd {
	ret := &Cmd{
		Cmd:    exec.CommandContext(ctx, name, arg...),
		logger: dlog.GetLogger(ctx),
	}
	ret.pidlock.Lock()
	return ret
}

func (c *Cmd) logiofn(prefix string) func(string) {
	return func(msg string) {
		c.pidlock.RLock()
		defer c.pidlock.RUnlock()
		pid := -1
		if c.Process != nil {
			pid = c.Process.Pid
		}
		c.logger.Printf("[pid:%v] %s %s", pid, prefix, msg)
	}
}

// Start starts the specified command but does not wait for it to complete.
//
// See the os/exec.Cmd.Start documenaton for more information.
func (c *Cmd) Start() error {
	c.Stdin = fixupReader(c.Stdin, c.logiofn("stdin  <"))
	if interfaceEqual(c.Stdout, c.Stderr) {
		c.Stdout = fixupWriter(c.Stdout, c.logiofn("stdout+stderr >"))
		c.Stderr = c.Stdout
	} else {
		c.Stdout = fixupWriter(c.Stdout, c.logiofn("stdout >"))
		c.Stderr = fixupWriter(c.Stderr, c.logiofn("stderr >"))
	}

	err := c.Cmd.Start()
	if err == nil {
		c.logger.Printf("[pid:%v] started command %#v", c.Process.Pid, c.Args)
		if stdin, isFile := c.Stdin.(*os.File); isFile {
			c.logger.Printf("[pid:%v] stdin  < not logging input read from file %s", c.Process.Pid, stdin.Name())
		}
		if stdout, isFile := c.Stdout.(*os.File); isFile {
			c.logger.Printf("[pid:%v] stdout > not logging output written to file %s", c.Process.Pid, stdout.Name())
		}
		if stderr, isFile := c.Stderr.(*os.File); isFile {
			c.logger.Printf("[pid:%v] stderr > not logging output written to file %s", c.Process.Pid, stderr.Name())
		}
	}
	c.pidlock.Unlock()
	return err
}

// Wait waits for the command to exit and waits for any copying to
// stdin or copying from stdout or stderr to complete.
//
// See the os/exec.Cmd.Wait documenaton for more information.
func (c *Cmd) Wait() error {
	err := c.Cmd.Wait()

	pid := -1
	if c.Process != nil {
		pid = c.Process.Pid
	}

	if err == nil {
		c.logger.Printf("[pid:%v] finished successfully: %v", pid, c.ProcessState)
	} else {
		c.logger.Printf("[pid:%v] finished with error: %v", pid, err)
	}

	return err
}

// StdinPipe returns a pipe that will be connected to the command's
// standard input when the command starts.
//
// This sets .Stdin to an *os.File, causing what you write to the pipe
// to not be logged.
//
// See the os/exec.Cmd.StdinPipe documenaton for more information.
func (c *Cmd) StdinPipe() (io.WriteCloser, error) { return c.Cmd.StdinPipe() }

// StdoutPipe returns a pipe that will be connected to the command's
// standard output when the command starts.
//
// This sets .Stdout to an *os.File, causing what you read from the
// pipe to not be logged.
//
// See the os/exec.Cmd.StdoutPipe documenaton for more information.
func (c *Cmd) StdoutPipe() (io.ReadCloser, error) { return c.Cmd.StdoutPipe() }

// StderrPipe returns a pipe that will be connected to the command's
// standard error when the command starts.
//
// This sets .Stderr to an *os.File, causing what you read from the
// pipe to not be logged.
//
// See the os/exec.Cmd.StderrPipe documenaton for more information.
func (c *Cmd) StderrPipe() (io.ReadCloser, error) { return c.Cmd.StderrPipe() }

// Higher-level methods around these implemented in borrowed_cmd.go
<commit_msg>(from AES) Outlaw dlog.Logger from appearing in any type or function signatures<commit_after>// Package dexec is a logging variant of os/exec.
//
// dexec is *almost* a drop-in replacement for os/exec.  Differences
// are:
//
// - The "Command" function is missing, because a context is always
// required; use CommandContext.
//
// - It is not valid to create a "Cmd" entirely by hand; you must
// create it using CommandContext.  After it has been created, you may
// adjust the fields as you would with an os/exec.Cmd.
//
// The logger used is configured in the context.Context passed to
// CommandContext by calling
// github.com/datawire/ambassador/pkg/dlog.WithLogger.
//
// A Cmd logs when it starts, its exit status, and if they aren't an
// *os.File, logs everything read from or written to .Stdin, .Stdout,
// and .Stderr.  If one of those is an *os.File (as it is following a
// call to .StdinPipe, .StdoutPipe, or .StderrPipe), then that stream
// won't be logged (but it will print a message at process-start
// noting that it isn't being logged).
//
// For example:
//
//     ctx := dlog.WithLogger(context.Background(), myLogger)
//     cmd := dexec.CommandContext(ctx, "printf", "%s\n", "foo bar", "baz")
//     cmd.Stdin = os.Stdin
//     err := cmd.Run()
//
// will log the lines
//
//     [pid:24272] started command []string{"printf", "%s\n", "foo bar", "baz"}
//     [pid:24272] stdin  < not logging input read from file /dev/stdin
//     [pid:24272] stdout+stderr > "foo bar\n"
//     [pid:24272] stdout+stderr > "baz\n"
//     [pid:24272] finished successfully: exit status 0
//
// If you would like a "pipe" to be logged, use an io.Pipe instead of
// calling .StdinPipe, .StdoutPipe, or .StderrPipe.
//
// See the os/exec documentation for more information.
package dexec

import (
	"context"
	"io"
	"os"
	"os/exec"
	"sync"

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

// Error is returned by LookPath when it fails to classify a file as an
// executable.
type Error = exec.Error

// An ExitError reports an unsuccessful exit by a command.
type ExitError = exec.ExitError

// ErrNotFound is the os/exec.ErrNotFound value.
var ErrNotFound = exec.ErrNotFound

// LookPath is the os/exe.LookPath function.
var LookPath = exec.LookPath

// Cmd represents an external command being prepared or run.
//
// A Cmd cannot be reused after calling its Run, Output or CombinedOutput
// methods.
//
// See the os/exec.Cmd documentation for information on the fields
// within it.
//
// Unlike an os/exec.Cmd, you MUST NOT construct a Cmd by hand, it
// must be created with CommandContext.
type Cmd struct {
	*exec.Cmd
	ctx context.Context

	pidlock sync.RWMutex
}

// CommandContext returns the Cmd struct to execute the named program with
// the given arguments.
//
// The provided context is used for two purposes:
//
//  1. To kill the process (by calling os.Process.Kill) if the context
//     becomes done before the command completes on its own.
//  2. To get the logger (by calling
//     github.com/datawire/ambassador/pkg/dlog.GetLogger on it).
//
// See the os/exec.Command and os/exec.CommandContext documentation
// for more information.
func CommandContext(ctx context.Context, name string, arg ...string) *Cmd {
	ret := &Cmd{
		Cmd: exec.CommandContext(ctx, name, arg...),
		ctx: ctx,
	}
	ret.pidlock.Lock()
	return ret
}

func (c *Cmd) logiofn(prefix string) func(string) {
	return func(msg string) {
		c.pidlock.RLock()
		defer c.pidlock.RUnlock()
		pid := -1
		if c.Process != nil {
			pid = c.Process.Pid
		}
		dlog.GetLogger(c.ctx).Printf("[pid:%v] %s %s", pid, prefix, msg)
	}
}

// Start starts the specified command but does not wait for it to complete.
//
// See the os/exec.Cmd.Start documenaton for more information.
func (c *Cmd) Start() error {
	c.Stdin = fixupReader(c.Stdin, c.logiofn("stdin  <"))
	if interfaceEqual(c.Stdout, c.Stderr) {
		c.Stdout = fixupWriter(c.Stdout, c.logiofn("stdout+stderr >"))
		c.Stderr = c.Stdout
	} else {
		c.Stdout = fixupWriter(c.Stdout, c.logiofn("stdout >"))
		c.Stderr = fixupWriter(c.Stderr, c.logiofn("stderr >"))
	}

	err := c.Cmd.Start()
	if err == nil {
		dlog.GetLogger(c.ctx).Printf("[pid:%v] started command %#v", c.Process.Pid, c.Args)
		if stdin, isFile := c.Stdin.(*os.File); isFile {
			dlog.GetLogger(c.ctx).Printf("[pid:%v] stdin  < not logging input read from file %s", c.Process.Pid, stdin.Name())
		}
		if stdout, isFile := c.Stdout.(*os.File); isFile {
			dlog.GetLogger(c.ctx).Printf("[pid:%v] stdout > not logging output written to file %s", c.Process.Pid, stdout.Name())
		}
		if stderr, isFile := c.Stderr.(*os.File); isFile {
			dlog.GetLogger(c.ctx).Printf("[pid:%v] stderr > not logging output written to file %s", c.Process.Pid, stderr.Name())
		}
	}
	c.pidlock.Unlock()
	return err
}

// Wait waits for the command to exit and waits for any copying to
// stdin or copying from stdout or stderr to complete.
//
// See the os/exec.Cmd.Wait documenaton for more information.
func (c *Cmd) Wait() error {
	err := c.Cmd.Wait()

	pid := -1
	if c.Process != nil {
		pid = c.Process.Pid
	}

	if err == nil {
		dlog.GetLogger(c.ctx).Printf("[pid:%v] finished successfully: %v", pid, c.ProcessState)
	} else {
		dlog.GetLogger(c.ctx).Printf("[pid:%v] finished with error: %v", pid, err)
	}

	return err
}

// StdinPipe returns a pipe that will be connected to the command's
// standard input when the command starts.
//
// This sets .Stdin to an *os.File, causing what you write to the pipe
// to not be logged.
//
// See the os/exec.Cmd.StdinPipe documenaton for more information.
func (c *Cmd) StdinPipe() (io.WriteCloser, error) { return c.Cmd.StdinPipe() }

// StdoutPipe returns a pipe that will be connected to the command's
// standard output when the command starts.
//
// This sets .Stdout to an *os.File, causing what you read from the
// pipe to not be logged.
//
// See the os/exec.Cmd.StdoutPipe documenaton for more information.
func (c *Cmd) StdoutPipe() (io.ReadCloser, error) { return c.Cmd.StdoutPipe() }

// StderrPipe returns a pipe that will be connected to the command's
// standard error when the command starts.
//
// This sets .Stderr to an *os.File, causing what you read from the
// pipe to not be logged.
//
// See the os/exec.Cmd.StderrPipe documenaton for more information.
func (c *Cmd) StderrPipe() (io.ReadCloser, error) { return c.Cmd.StderrPipe() }

// Higher-level methods around these implemented in borrowed_cmd.go
<|endoftext|>