<commit_before>package loaders

import (
	"regexp"
	"strings"

	_ "github.com/mattn/go-sqlite3"

	"github.com/knq/xo/internal"
	"github.com/knq/xo/models"
)

func init() {
	internal.SchemaLoaders["sqlite3"] = internal.TypeLoader{
		ProcessRelkind: SqRelkind,
		ParamN:         func(int) string { return "?" },
		MaskFunc:       func() string { return "?" },
		ParseType:      SqParseType,
		TableList: func(db models.XODB, schema string, relkind string) ([]*models.Table, error) {
			return models.SqTables(db, relkind)
		},
		ColumnList: func(db models.XODB, schema string, table string) ([]*models.Column, error) {
			return models.SqTableColumns(db, table)
		},
		ForeignKeyList: func(db models.XODB, schema string, table string) ([]*models.ForeignKey, error) {
			return models.SqTableForeignKeys(db, table)
		},
		IndexList: func(db models.XODB, schema string, table string) ([]*models.Index, error) {
			return models.SqTableIndexes(db, table)
		},
		IndexColumnList: func(db models.XODB, schema string, table string, index string) ([]*models.IndexColumn, error) {
			return models.SqIndexColumns(db, index)
		},
		QueryColumnList: SqQueryColumns,
	}
}

// SqRelkind returns the sqlite string representation for RelType.
func SqRelkind(relType internal.RelType) string {
	var s string
	switch relType {
	case internal.Table:
		s = "table"
	case internal.View:
		s = "view"
	default:
		panic("unsupported RelType")
	}
	return s
}

var uRE = regexp.MustCompile(`\s*unsigned\*`)

// SqParseType parse a postgres type into a Go type based on the column
// definition.
func SqParseType(args *internal.ArgType, dt string, nullable bool) (int, string, string) {
	precision := 0
	nilVal := "nil"
	unsigned := false

	// extract precision
	dt, precision, _ = args.ParsePrecision(dt)

	if uRE.MatchString(dt) {
		unsigned = true
		uRE.ReplaceAllString(dt, "")
	}

	var typ string
	switch dt {
	case "bool", "boolean":
		nilVal = "false"
		typ = "bool"
		if nullable {
			nilVal = "sql.NullBool{}"
			typ = "sql.NullBool"
		}

	case "int", "integer", "tinyint", "smallint", "mediumint", "bigint":
		nilVal = "0"
		typ = args.Int32Type
		if nullable {
			nilVal = "sql.NullInt64{}"
			typ = "sql.NullInt64"
		}

	case "numeric", "real", "double", "float", "decimal":
		nilVal = "0.0"
		typ = "float64"
		if nullable {
			nilVal = "sql.NullFloat64{}"
			typ = "sql.NullFloat64"
		}

	case "blob":
		typ = "[]byte"

	default:
		// case "varchar", "character", "varying character", "nchar", "native character", "nvarchar", "text", "clob", "datetime", "date", "time":
		nilVal = `""`
		typ = "string"
		if nullable {
			nilVal = "sql.NullString{}"
			typ = "sql.NullString"
		}
	}

	// if unsigned ...
	if internal.IntRE.MatchString(typ) && unsigned {
		typ = "u" + typ
	}

	return precision, nilVal, typ
}

// SqQueryColumns parses a sqlite query and generates a type for it.
func SqQueryColumns(args *internal.ArgType, inspect []string) ([]*models.Column, error) {
	var err error

	// create temporary view xoid
	xoid := "_xo_" + internal.GenRandomID()
	viewq := `CREATE TEMPORARY VIEW ` + xoid + ` AS ` + strings.Join(inspect, "\n")
	models.XOLog(viewq)
	_, err = args.DB.Exec(viewq)
	if err != nil {
		return nil, err
	}

	// load column information
	return models.SqTableColumns(args.DB, xoid)
}
<commit_msg>Fixing case issue with SQLite databases<commit_after>package loaders

import (
	"regexp"
	"strings"

	_ "github.com/mattn/go-sqlite3"

	"github.com/knq/xo/internal"
	"github.com/knq/xo/models"
)

func init() {
	internal.SchemaLoaders["sqlite3"] = internal.TypeLoader{
		ProcessRelkind: SqRelkind,
		ParamN:         func(int) string { return "?" },
		MaskFunc:       func() string { return "?" },
		ParseType:      SqParseType,
		TableList: func(db models.XODB, schema string, relkind string) ([]*models.Table, error) {
			return models.SqTables(db, relkind)
		},
		ColumnList: func(db models.XODB, schema string, table string) ([]*models.Column, error) {
			return models.SqTableColumns(db, table)
		},
		ForeignKeyList: func(db models.XODB, schema string, table string) ([]*models.ForeignKey, error) {
			return models.SqTableForeignKeys(db, table)
		},
		IndexList: func(db models.XODB, schema string, table string) ([]*models.Index, error) {
			return models.SqTableIndexes(db, table)
		},
		IndexColumnList: func(db models.XODB, schema string, table string, index string) ([]*models.IndexColumn, error) {
			return models.SqIndexColumns(db, index)
		},
		QueryColumnList: SqQueryColumns,
	}
}

// SqRelkind returns the sqlite string representation for RelType.
func SqRelkind(relType internal.RelType) string {
	var s string
	switch relType {
	case internal.Table:
		s = "table"
	case internal.View:
		s = "view"
	default:
		panic("unsupported RelType")
	}
	return s
}

var uRE = regexp.MustCompile(`\s*unsigned\*`)

// SqParseType parse a postgres type into a Go type based on the column
// definition.
func SqParseType(args *internal.ArgType, dt string, nullable bool) (int, string, string) {
	precision := 0
	nilVal := "nil"
	unsigned := false

	dt = strings.ToLower(dt)

	// extract precision
	dt, precision, _ = args.ParsePrecision(dt)

	if uRE.MatchString(dt) {
		unsigned = true
		uRE.ReplaceAllString(dt, "")
	}

	var typ string
	switch dt {
	case "bool", "boolean":
		nilVal = "false"
		typ = "bool"
		if nullable {
			nilVal = "sql.NullBool{}"
			typ = "sql.NullBool"
		}

	case "int", "integer", "tinyint", "smallint", "mediumint", "bigint":
		nilVal = "0"
		typ = args.Int32Type
		if nullable {
			nilVal = "sql.NullInt64{}"
			typ = "sql.NullInt64"
		}

	case "numeric", "real", "double", "float", "decimal":
		nilVal = "0.0"
		typ = "float64"
		if nullable {
			nilVal = "sql.NullFloat64{}"
			typ = "sql.NullFloat64"
		}

	case "blob":
		typ = "[]byte"

	default:
		// case "varchar", "character", "varying character", "nchar", "native character", "nvarchar", "text", "clob", "datetime", "date", "time":
		nilVal = `""`
		typ = "string"
		if nullable {
			nilVal = "sql.NullString{}"
			typ = "sql.NullString"
		}
	}

	// if unsigned ...
	if internal.IntRE.MatchString(typ) && unsigned {
		typ = "u" + typ
	}

	return precision, nilVal, typ
}

// SqQueryColumns parses a sqlite query and generates a type for it.
func SqQueryColumns(args *internal.ArgType, inspect []string) ([]*models.Column, error) {
	var err error

	// create temporary view xoid
	xoid := "_xo_" + internal.GenRandomID()
	viewq := `CREATE TEMPORARY VIEW ` + xoid + ` AS ` + strings.Join(inspect, "\n")
	models.XOLog(viewq)
	_, err = args.DB.Exec(viewq)
	if err != nil {
		return nil, err
	}

	// load column information
	return models.SqTableColumns(args.DB, xoid)
}
<|endoftext|>


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

package sql

import (
	"fmt"
	"time"

	"github.com/cockroachdb/cockroach/client"
	"github.com/cockroachdb/cockroach/roachpb"
	"github.com/cockroachdb/cockroach/sql/parser"
	"github.com/opentracing/basictracer-go"
	"github.com/opentracing/opentracing-go"
)

// explainTraceNode is a planNode that wraps another node and converts DebugValues() results to a
// row of Values(). It is used as the top-level node for EXPLAIN (TRACE) statements.
type explainTraceNode struct {
	plan planNode
	txn  *client.Txn
	// Internal state, not to be initialized.
	exhausted bool
	rows      []parser.DTuple
	lastTS    time.Time
	lastPos   int
}

var traceColumns = append([]ResultColumn{
	{Name: "Cumulative Time", Typ: parser.DummyTimestamp},
	{Name: "Duration", Typ: parser.DummyString},
	{Name: "Span Pos", Typ: parser.DummyInt},
	{Name: "Operation", Typ: parser.DummyString},
	{Name: "Event", Typ: parser.DummyString},
}, debugColumns...)

func (*explainTraceNode) Columns() []ResultColumn { return traceColumns }
func (*explainTraceNode) Ordering() orderingInfo  { return orderingInfo{} }

func (n *explainTraceNode) PErr() *roachpb.Error { return nil }
func (n *explainTraceNode) Next() bool {
	first := n.rows == nil
	if first {
		n.rows = []parser.DTuple{}
	}
	for !n.exhausted && len(n.rows) <= 1 {
		var vals debugValues
		if !n.plan.Next() {
			n.exhausted = true
			if pErr := n.PErr(); pErr != nil {
				n.txn.Trace.LogEvent(pErr.GoError().Error())
			}
			n.txn.Trace.LogEvent("tracing completed")
			n.txn.Trace.Finish()
			n.txn.Trace = nil
		} else {
			vals = n.plan.DebugValues()
		}
		var basePos int
		if len(n.txn.CollectedSpans) == 0 {
			if !n.exhausted {
				n.txn.CollectedSpans = append(n.txn.CollectedSpans, basictracer.RawSpan{
					Context: basictracer.Context{},
					Logs:    []opentracing.LogData{{Timestamp: n.lastTS}},
				})
			}
			basePos = n.lastPos + 1
		}

		// Iterate through once to determine earliest timestamp.
		var earliest time.Time
		for _, sp := range n.txn.CollectedSpans {
			for _, entry := range sp.Logs {
				if earliest.IsZero() || entry.Timestamp.Before(earliest) {
					earliest = entry.Timestamp
				}
			}
		}

		for _, sp := range n.txn.CollectedSpans {
			for i, entry := range sp.Logs {
				commulativeDuration := fmt.Sprintf("%.3fms", time.Duration(entry.Timestamp.Sub(earliest)).Seconds()*1000)
				var duration string
				if i > 0 {
					duration = fmt.Sprintf("%.3fms", time.Duration(entry.Timestamp.Sub(n.lastTS)).Seconds()*1000)
				}
				cols := append(parser.DTuple{
					parser.DString(commulativeDuration),
					parser.DString(duration),
					parser.DInt(basePos + i),
					parser.DString(sp.Operation),
					parser.DString(entry.Event),
				}, vals.AsRow()...)

				// Timestamp is added for sorting, but will be removed after sort.
				n.rows = append(n.rows, append(cols, parser.DTimestamp{Time: entry.Timestamp}))
				n.lastTS, n.lastPos = entry.Timestamp, i
			}
		}
		n.txn.CollectedSpans = nil
	}

	if first {
		return len(n.rows) > 0
	}
	if len(n.rows) <= 1 {
		return false
	}
	n.rows = n.rows[1:]
	return true
}

func (n *explainTraceNode) ExplainPlan() (name, description string, children []planNode) {
	return n.plan.ExplainPlan()
}

func (n *explainTraceNode) Values() parser.DTuple {
	return n.rows[0]
}

func (*explainTraceNode) DebugValues() debugValues {
	panic("debug mode not implemented in explainTraceNode")
}

func (*explainTraceNode) SetLimitHint(_ int64) {}
<commit_msg>sql: fix Explain mis-matched col type test break<commit_after>// Copyright 2016 The Cockroach Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
// implied. See the License for the specific language governing
// permissions and limitations under the License.
//
// Author: Tobias Schottdorf (tobias.schottdorf@gmail.com)

package sql

import (
	"fmt"
	"time"

	"github.com/cockroachdb/cockroach/client"
	"github.com/cockroachdb/cockroach/roachpb"
	"github.com/cockroachdb/cockroach/sql/parser"
	"github.com/opentracing/basictracer-go"
	"github.com/opentracing/opentracing-go"
)

// explainTraceNode is a planNode that wraps another node and converts DebugValues() results to a
// row of Values(). It is used as the top-level node for EXPLAIN (TRACE) statements.
type explainTraceNode struct {
	plan planNode
	txn  *client.Txn
	// Internal state, not to be initialized.
	exhausted bool
	rows      []parser.DTuple
	lastTS    time.Time
	lastPos   int
}

var traceColumns = append([]ResultColumn{
	{Name: "Cumulative Time", Typ: parser.DummyString},
	{Name: "Duration", Typ: parser.DummyString},
	{Name: "Span Pos", Typ: parser.DummyInt},
	{Name: "Operation", Typ: parser.DummyString},
	{Name: "Event", Typ: parser.DummyString},
}, debugColumns...)

func (*explainTraceNode) Columns() []ResultColumn { return traceColumns }
func (*explainTraceNode) Ordering() orderingInfo  { return orderingInfo{} }

func (n *explainTraceNode) PErr() *roachpb.Error { return nil }
func (n *explainTraceNode) Next() bool {
	first := n.rows == nil
	if first {
		n.rows = []parser.DTuple{}
	}
	for !n.exhausted && len(n.rows) <= 1 {
		var vals debugValues
		if !n.plan.Next() {
			n.exhausted = true
			if pErr := n.PErr(); pErr != nil {
				n.txn.Trace.LogEvent(pErr.GoError().Error())
			}
			n.txn.Trace.LogEvent("tracing completed")
			n.txn.Trace.Finish()
			n.txn.Trace = nil
		} else {
			vals = n.plan.DebugValues()
		}
		var basePos int
		if len(n.txn.CollectedSpans) == 0 {
			if !n.exhausted {
				n.txn.CollectedSpans = append(n.txn.CollectedSpans, basictracer.RawSpan{
					Context: basictracer.Context{},
					Logs:    []opentracing.LogData{{Timestamp: n.lastTS}},
				})
			}
			basePos = n.lastPos + 1
		}

		// Iterate through once to determine earliest timestamp.
		var earliest time.Time
		for _, sp := range n.txn.CollectedSpans {
			for _, entry := range sp.Logs {
				if earliest.IsZero() || entry.Timestamp.Before(earliest) {
					earliest = entry.Timestamp
				}
			}
		}

		for _, sp := range n.txn.CollectedSpans {
			for i, entry := range sp.Logs {
				commulativeDuration := fmt.Sprintf("%.3fms", time.Duration(entry.Timestamp.Sub(earliest)).Seconds()*1000)
				var duration string
				if i > 0 {
					duration = fmt.Sprintf("%.3fms", time.Duration(entry.Timestamp.Sub(n.lastTS)).Seconds()*1000)
				}
				cols := append(parser.DTuple{
					parser.DString(commulativeDuration),
					parser.DString(duration),
					parser.DInt(basePos + i),
					parser.DString(sp.Operation),
					parser.DString(entry.Event),
				}, vals.AsRow()...)

				// Timestamp is added for sorting, but will be removed after sort.
				n.rows = append(n.rows, append(cols, parser.DTimestamp{Time: entry.Timestamp}))
				n.lastTS, n.lastPos = entry.Timestamp, i
			}
		}
		n.txn.CollectedSpans = nil
	}

	if first {
		return len(n.rows) > 0
	}
	if len(n.rows) <= 1 {
		return false
	}
	n.rows = n.rows[1:]
	return true
}

func (n *explainTraceNode) ExplainPlan() (name, description string, children []planNode) {
	return n.plan.ExplainPlan()
}

func (n *explainTraceNode) Values() parser.DTuple {
	return n.rows[0]
}

func (*explainTraceNode) DebugValues() debugValues {
	panic("debug mode not implemented in explainTraceNode")
}

func (*explainTraceNode) SetLimitHint(_ int64) {}
<|endoftext|>


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

package main

import (
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"runtime"
	"strings"

	"github.com/GoogleCloudPlatform/guest-logging-go/logger"
)

var (
	googleComment    = "# Added by Google Compute Engine OS Login."
	googleBlockStart = "#### Google OS Login control. Do not edit this section. ####"
	googleBlockEnd   = "#### End Google OS Login control section. ####"
)

type osloginMgr struct{}

// We also read project keys first, letting instance-level keys take
// precedence.
func getOSLoginEnabled(md *metadata) (bool, bool) {
	var enable bool
	if md.Project.Attributes.EnableOSLogin != nil {
		enable = *md.Project.Attributes.EnableOSLogin
	}
	if md.Instance.Attributes.EnableOSLogin != nil {
		enable = *md.Instance.Attributes.EnableOSLogin
	}
	var twofactor bool
	if md.Project.Attributes.TwoFactor != nil {
		twofactor = *md.Project.Attributes.TwoFactor
	}
	if md.Instance.Attributes.TwoFactor != nil {
		twofactor = *md.Instance.Attributes.TwoFactor
	}
	return enable, twofactor
}

func (o *osloginMgr) diff() bool {
	oldEnable, oldTwoFactor := getOSLoginEnabled(oldMetadata)
	enable, twofactor := getOSLoginEnabled(newMetadata)
	return oldMetadata.Project.ProjectID == "" ||
		// True on first run or if any value has changed.
		(oldTwoFactor != twofactor) ||
		(oldEnable != enable)
}

func (o *osloginMgr) timeout() bool {
	return false
}

func (o *osloginMgr) disabled(os string) bool {
	return os == "windows"
}

func (o *osloginMgr) set() error {
	// We need to know if it was previously enabled for the clearing of
	// metadata-based SSH keys.
	oldEnable, _ := getOSLoginEnabled(oldMetadata)
	enable, twofactor := getOSLoginEnabled(newMetadata)

	// Proxy for detecting whether we think it is currently enabled. We
	// read files instead of any variables for this because we want to
	// handle agent being down during transitions.
	var isEnabled bool
	if nsswitch, err := ioutil.ReadFile("/etc/nsswitch.conf"); err == nil && strings.Contains(string(nsswitch), "oslogin") {
		isEnabled = true
	}

	if enable && !oldEnable {
		logger.Infof("Enabling OS Login")
		newMetadata.Instance.Attributes.SSHKeys = nil
		newMetadata.Project.Attributes.SSHKeys = nil
		(&accountsMgr{}).set()
	}

	if err := writeSSHConfig(enable, twofactor); err != nil {
		logger.Errorf("Error updating SSH config: %v.", err)
	}

	if err := writeNSSwitchConfig(enable); err != nil {
		logger.Errorf("Error updating NSS config: %v.", err)
	}

	if err := writePAMConfig(enable, twofactor); err != nil {
		logger.Errorf("Error updating PAM config: %v.", err)
	}

	if enable {
		if err := createOSLoginDirs(); err != nil {
			logger.Errorf("Error creating OS Login directory: %v.", err)
		}

		if err := createOSLoginSudoersFile(); err != nil {
			logger.Errorf("Error creating OS Login sudoers file: %v.", err)
		}

		// Services which need to be restarted primarily due to caching
		// issues or config changes. Only do this if we think OS Login
		// was not already enabled.
		if !isEnabled {
			for _, svc := range []string{"ssh", "sshd", "nscd", "unscd", "systemd-logind", "cron", "crond"} {
				if err := restartService(svc); err != nil {
					logger.Errorf("Error restarting service: %v.", err)
				}
			}
		}

		if err := runCmd(exec.Command("google_oslogin_nss_cache")); err != nil {
			logger.Errorf("Error updating NSS cache: %v.", err)
		}
	}

	return nil
}

func filterGoogleLines(contents string) []string {
	var isgoogle, isgoogleblock bool
	var filtered []string
	for _, line := range strings.Split(contents, "\n") {
		switch {
		case strings.Contains(line, googleComment) && !isgoogleblock:
			isgoogle = true
		case strings.Contains(line, googleBlockEnd):
			isgoogleblock = false
			isgoogle = false
		case isgoogleblock, strings.Contains(line, googleBlockStart):
			isgoogleblock = true
		case isgoogle:
			isgoogle = false
		default:
			filtered = append(filtered, line)
		}
	}
	return filtered
}

func writeConfigFile(path, contents string) error {
	file, err := os.OpenFile(path, os.O_WRONLY|os.O_TRUNC, 0777)
	if err != nil {
		return err
	}
	defer closeFile(file)
	file.WriteString(contents)
	return nil
}

func updateSSHConfig(sshConfig string, enable, twofactor bool) string {
	// TODO: this feels like a case for a text/template
	challengeResponseEnable := "ChallengeResponseAuthentication yes"
	authorizedKeysCommand := "AuthorizedKeysCommand /usr/bin/google_authorized_keys"
	if runtime.GOOS == "freebsd" {
		authorizedKeysCommand = "AuthorizedKeysCommand /usr/local/bin/google_authorized_keys"
	}
	authorizedKeysUser := "AuthorizedKeysCommandUser root"
	twoFactorAuthMethods := "AuthenticationMethods publickey,keyboard-interactive"
	if (osRelease.os == "rhel" || osRelease.os == "centos") && osRelease.version.major == 6 {
		authorizedKeysUser = "AuthorizedKeysCommandRunAs root"
		twoFactorAuthMethods = "RequiredAuthentications2 publickey,keyboard-interactive"
	}

	filtered := filterGoogleLines(string(sshConfig))

	if enable {
		osLoginBlock := []string{googleBlockStart, authorizedKeysCommand, authorizedKeysUser}
		if twofactor {
			osLoginBlock = append(osLoginBlock, twoFactorAuthMethods, challengeResponseEnable)
		}
		osLoginBlock = append(osLoginBlock, googleBlockEnd)
		filtered = append(osLoginBlock, filtered...)
	}

	return strings.Join(filtered, "\n")
}

func writeSSHConfig(enable, twofactor bool) error {
	sshConfig, err := ioutil.ReadFile("/etc/ssh/sshd_config")
	if err != nil {
		return err
	}
	proposed := updateSSHConfig(string(sshConfig), enable, twofactor)
	if proposed == string(sshConfig) {
		return nil
	}
	return writeConfigFile("/etc/ssh/sshd_config", proposed)
}

func updateNSSwitchConfig(nsswitch string, enable bool) string {
	oslogin := " cache_oslogin oslogin"

	var filtered []string
	for _, line := range strings.Split(string(nsswitch), "\n") {
		if strings.HasPrefix(line, "passwd:") || strings.HasPrefix(line, "group:") {
			present := strings.Contains(line, "oslogin")
			if enable && !present {
				line += oslogin
			} else if !enable && present {
				line = strings.TrimSuffix(line, oslogin)
			}

			if runtime.GOOS == "freebsd" {
				line = strings.Replace(line, "compat", "files", 1)
			}
		}
		filtered = append(filtered, line)
	}
	return strings.Join(filtered, "\n")
}

func writeNSSwitchConfig(enable bool) error {
	nsswitch, err := ioutil.ReadFile("/etc/nsswitch.conf")
	if err != nil {
		return err
	}
	proposed := updateNSSwitchConfig(string(nsswitch), enable)
	if proposed == string(nsswitch) {
		return nil
	}
	return writeConfigFile("/etc/nsswitch.conf", proposed)
}

// Adds entries to the PAM config for sshd and su which reflect the current
// enablements. Only writes files if they have changed from what's on disk.
func updatePAMsshd(pamsshd string, enable, twofactor bool) string {
	authOSLogin := "auth       [success=done perm_denied=die default=ignore] pam_oslogin_login.so"
	authGroup := "auth       [default=ignore] pam_group.so"
	accountOSLogin := "account    [success=ok ignore=ignore default=die] pam_oslogin_login.so"
	accountOSLoginAdmin := "account    [success=ok default=ignore] pam_oslogin_admin.so"
	sessionHomeDir := "session    [success=ok default=ignore] pam_mkhomedir.so"

	if runtime.GOOS == "freebsd" {
		authOSLogin = "auth       optional pam_oslogin_login.so"
		authGroup = "auth       optional pam_group.so"
		accountOSLogin = "account    requisite pam_oslogin_login.so"
		accountOSLoginAdmin = "account    optional pam_oslogin_admin.so"
		sessionHomeDir = "session    optional pam_mkhomedir.so"
	}

	filtered := filterGoogleLines(string(pamsshd))
	if enable {
		topOfFile := []string{googleBlockStart}
		if twofactor {
			topOfFile = append(topOfFile, authOSLogin)
		}
		topOfFile = append(topOfFile, authGroup, googleBlockEnd)
		bottomOfFile := []string{googleBlockStart, accountOSLogin, accountOSLoginAdmin, sessionHomeDir, googleBlockEnd}
		filtered = append(topOfFile, filtered...)
		filtered = append(filtered, bottomOfFile...)
	}
	return strings.Join(filtered, "\n")
}

func updatePAMsu(pamsu string, enable bool) string {
	accountSu := "account    [success=bad ignore=ignore] pam_oslogin_login.so"

	filtered := filterGoogleLines(pamsu)
	if enable {
		filtered = append([]string{googleComment, accountSu}, filtered...)
	}

	return strings.Join(filtered, "\n")
}

func writePAMConfig(enable, twofactor bool) error {
	pamsshd, err := ioutil.ReadFile("/etc/pam.d/sshd")
	if err != nil {
		return err
	}
	proposed := updatePAMsshd(string(pamsshd), enable, twofactor)
	if proposed != string(pamsshd) {
		if err := writeConfigFile("/etc/pam.d/sshd", proposed); err != nil {
			return err
		}
	}

	pamsu, err := ioutil.ReadFile("/etc/pam.d/su")
	if err != nil {
		return err
	}
	proposed = updatePAMsu(string(pamsu), enable)
	if proposed != string(pamsu) {
		if err := writeConfigFile("/etc/pam.d/su", proposed); err != nil {
			return err
		}
	}

	return nil
}

// Creates necessary OS Login directories if they don't exist.
func createOSLoginDirs() error {
	restorecon, restoreconerr := exec.LookPath("restorecon")

	for _, dir := range []string{"/var/google-sudoers.d", "/var/google-users.d"} {
		err := os.Mkdir(dir, 0750)
		if err != nil && !os.IsExist(err) {
			return err
		}
		if restoreconerr == nil {
			runCmd(exec.Command(restorecon, dir))
		}
	}
	return nil
}

func createOSLoginSudoersFile() error {
	osloginSudoers := "/etc/sudoers.d/google-oslogin"
	if runtime.GOOS == "freebsd" {
		osloginSudoers = "/usr/local" + osloginSudoers
	}
	sudoFile, err := os.OpenFile(osloginSudoers, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0440)
	if err != nil {
		if os.IsExist(err) {
			return nil
		}
		return err
	}
	fmt.Fprintf(sudoFile, "#includedir /var/google-sudoers.d\n")
	return sudoFile.Close()
}

// restartService tries to restart a service on linux-like systems. It attempts
// to find and use the following mechanisms in order:
// 1. The `systemctl` utility, if in a systemd environment.
// 2. The `service` command, if present.
// 3. A SysVinit script directly, if present.
// Missing mechanisms and missing or disabled services are ignored.
func restartService(servicename string) error {
	init, err := os.Readlink("/sbin/init")
	if err == nil && strings.Contains(init, "systemd") {
		if systemctl, err := exec.LookPath("systemctl"); err == nil {
			if err := runCmd(exec.Command(systemctl, "is-active", servicename+".service")); err == nil {
				return runCmd(exec.Command(systemctl, "restart", servicename+".service"))
			}
			return nil
		}
	}
	service, err := exec.LookPath("service")
	if err == nil {
		if err := runCmd(exec.Command(service, servicename, "status")); err == nil {
			return runCmd(exec.Command(service, servicename, "restart"))
		}
		return nil
	}
	initService := "/etc/init.d/" + servicename
	if _, err := os.Stat(initService); err == nil {
		if err := runCmd(exec.Command(initService, "status")); err == nil {
			return runCmd(exec.Command(initService, "restart"))
		}
		return nil
	}

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

package main

import (
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"runtime"
	"strings"

	"github.com/GoogleCloudPlatform/guest-logging-go/logger"
)

var (
	googleComment    = "# Added by Google Compute Engine OS Login."
	googleBlockStart = "#### Google OS Login control. Do not edit this section. ####"
	googleBlockEnd   = "#### End Google OS Login control section. ####"
)

type osloginMgr struct{}

// We also read project keys first, letting instance-level keys take
// precedence.
func getOSLoginEnabled(md *metadata) (bool, bool) {
	var enable bool
	if md.Project.Attributes.EnableOSLogin != nil {
		enable = *md.Project.Attributes.EnableOSLogin
	}
	if md.Instance.Attributes.EnableOSLogin != nil {
		enable = *md.Instance.Attributes.EnableOSLogin
	}
	var twofactor bool
	if md.Project.Attributes.TwoFactor != nil {
		twofactor = *md.Project.Attributes.TwoFactor
	}
	if md.Instance.Attributes.TwoFactor != nil {
		twofactor = *md.Instance.Attributes.TwoFactor
	}
	return enable, twofactor
}

func (o *osloginMgr) diff() bool {
	oldEnable, oldTwoFactor := getOSLoginEnabled(oldMetadata)
	enable, twofactor := getOSLoginEnabled(newMetadata)
	return oldMetadata.Project.ProjectID == "" ||
		// True on first run or if any value has changed.
		(oldTwoFactor != twofactor) ||
		(oldEnable != enable)
}

func (o *osloginMgr) timeout() bool {
	return false
}

func (o *osloginMgr) disabled(os string) bool {
	return os == "windows"
}

func (o *osloginMgr) set() error {
	// We need to know if it was previously enabled for the clearing of
	// metadata-based SSH keys.
	oldEnable, _ := getOSLoginEnabled(oldMetadata)
	enable, twofactor := getOSLoginEnabled(newMetadata)

	if enable && !oldEnable {
		logger.Infof("Enabling OS Login")
		newMetadata.Instance.Attributes.SSHKeys = nil
		newMetadata.Project.Attributes.SSHKeys = nil
		(&accountsMgr{}).set()
	}

	if err := writeSSHConfig(enable, twofactor); err != nil {
		logger.Errorf("Error updating SSH config: %v.", err)
	}

	if err := writeNSSwitchConfig(enable); err != nil {
		logger.Errorf("Error updating NSS config: %v.", err)
	}

	if err := writePAMConfig(enable, twofactor); err != nil {
		logger.Errorf("Error updating PAM config: %v.", err)
	}

	for _, svc := range []string{"ssh", "sshd", "nscd", "unscd", "systemd-logind", "cron", "crond"} {
		if err := restartService(svc); err != nil {
			logger.Errorf("Error restarting service: %v.", err)
		}
	}

	if enable {
		if err := createOSLoginDirs(); err != nil {
			logger.Errorf("Error creating OS Login directory: %v.", err)
		}

		if err := createOSLoginSudoersFile(); err != nil {
			logger.Errorf("Error creating OS Login sudoers file: %v.", err)
		}

		if err := runCmd(exec.Command("google_oslogin_nss_cache")); err != nil {
			logger.Errorf("Error updating NSS cache: %v.", err)
		}

	}

	return nil
}

func filterGoogleLines(contents string) []string {
	var isgoogle, isgoogleblock bool
	var filtered []string
	for _, line := range strings.Split(contents, "\n") {
		switch {
		case strings.Contains(line, googleComment) && !isgoogleblock:
			isgoogle = true
		case strings.Contains(line, googleBlockEnd):
			isgoogleblock = false
			isgoogle = false
		case isgoogleblock, strings.Contains(line, googleBlockStart):
			isgoogleblock = true
		case isgoogle:
			isgoogle = false
		default:
			filtered = append(filtered, line)
		}
	}
	return filtered
}

func writeConfigFile(path, contents string) error {
	file, err := os.OpenFile(path, os.O_WRONLY|os.O_TRUNC, 0777)
	if err != nil {
		return err
	}
	defer closeFile(file)
	file.WriteString(contents)
	return nil
}

func updateSSHConfig(sshConfig string, enable, twofactor bool) string {
	// TODO: this feels like a case for a text/template
	challengeResponseEnable := "ChallengeResponseAuthentication yes"
	authorizedKeysCommand := "AuthorizedKeysCommand /usr/bin/google_authorized_keys"
	if runtime.GOOS == "freebsd" {
		authorizedKeysCommand = "AuthorizedKeysCommand /usr/local/bin/google_authorized_keys"
	}
	authorizedKeysUser := "AuthorizedKeysCommandUser root"
	twoFactorAuthMethods := "AuthenticationMethods publickey,keyboard-interactive"
	if (osRelease.os == "rhel" || osRelease.os == "centos") && osRelease.version.major == 6 {
		authorizedKeysUser = "AuthorizedKeysCommandRunAs root"
		twoFactorAuthMethods = "RequiredAuthentications2 publickey,keyboard-interactive"
	}

	filtered := filterGoogleLines(string(sshConfig))

	if enable {
		osLoginBlock := []string{googleBlockStart, authorizedKeysCommand, authorizedKeysUser}
		if twofactor {
			osLoginBlock = append(osLoginBlock, twoFactorAuthMethods, challengeResponseEnable)
		}
		osLoginBlock = append(osLoginBlock, googleBlockEnd)
		filtered = append(osLoginBlock, filtered...)
	}

	return strings.Join(filtered, "\n")
}

func writeSSHConfig(enable, twofactor bool) error {
	sshConfig, err := ioutil.ReadFile("/etc/ssh/sshd_config")
	if err != nil {
		return err
	}
	proposed := updateSSHConfig(string(sshConfig), enable, twofactor)
	if proposed == string(sshConfig) {
		return nil
	}
	return writeConfigFile("/etc/ssh/sshd_config", proposed)
}

func updateNSSwitchConfig(nsswitch string, enable bool) string {
	oslogin := " cache_oslogin oslogin"

	var filtered []string
	for _, line := range strings.Split(string(nsswitch), "\n") {
		if strings.HasPrefix(line, "passwd:") || strings.HasPrefix(line, "group:") {
			present := strings.Contains(line, "oslogin")
			if enable && !present {
				line += oslogin
			} else if !enable && present {
				line = strings.TrimSuffix(line, oslogin)
			}

			if runtime.GOOS == "freebsd" {
				line = strings.Replace(line, "compat", "files", 1)
			}
		}
		filtered = append(filtered, line)
	}
	return strings.Join(filtered, "\n")
}

func writeNSSwitchConfig(enable bool) error {
	nsswitch, err := ioutil.ReadFile("/etc/nsswitch.conf")
	if err != nil {
		return err
	}
	proposed := updateNSSwitchConfig(string(nsswitch), enable)
	if proposed == string(nsswitch) {
		return nil
	}
	return writeConfigFile("/etc/nsswitch.conf", proposed)
}

// Adds entries to the PAM config for sshd and su which reflect the current
// enablements. Only writes files if they have changed from what's on disk.
func updatePAMsshd(pamsshd string, enable, twofactor bool) string {
	authOSLogin := "auth       [success=done perm_denied=die default=ignore] pam_oslogin_login.so"
	authGroup := "auth       [default=ignore] pam_group.so"
	accountOSLogin := "account    [success=ok ignore=ignore default=die] pam_oslogin_login.so"
	accountOSLoginAdmin := "account    [success=ok default=ignore] pam_oslogin_admin.so"
	sessionHomeDir := "session    [success=ok default=ignore] pam_mkhomedir.so"

	if runtime.GOOS == "freebsd" {
		authOSLogin = "auth       optional pam_oslogin_login.so"
		authGroup = "auth       optional pam_group.so"
		accountOSLogin = "account    requisite pam_oslogin_login.so"
		accountOSLoginAdmin = "account    optional pam_oslogin_admin.so"
		sessionHomeDir = "session    optional pam_mkhomedir.so"
	}

	filtered := filterGoogleLines(string(pamsshd))
	if enable {
		topOfFile := []string{googleBlockStart}
		if twofactor {
			topOfFile = append(topOfFile, authOSLogin)
		}
		topOfFile = append(topOfFile, authGroup, googleBlockEnd)
		bottomOfFile := []string{googleBlockStart, accountOSLogin, accountOSLoginAdmin, sessionHomeDir, googleBlockEnd}
		filtered = append(topOfFile, filtered...)
		filtered = append(filtered, bottomOfFile...)
	}
	return strings.Join(filtered, "\n")
}

func updatePAMsu(pamsu string, enable bool) string {
	accountSu := "account    [success=bad ignore=ignore] pam_oslogin_login.so"

	filtered := filterGoogleLines(pamsu)
	if enable {
		filtered = append([]string{googleComment, accountSu}, filtered...)
	}

	return strings.Join(filtered, "\n")
}

func writePAMConfig(enable, twofactor bool) error {
	pamsshd, err := ioutil.ReadFile("/etc/pam.d/sshd")
	if err != nil {
		return err
	}
	proposed := updatePAMsshd(string(pamsshd), enable, twofactor)
	if proposed != string(pamsshd) {
		if err := writeConfigFile("/etc/pam.d/sshd", proposed); err != nil {
			return err
		}
	}

	pamsu, err := ioutil.ReadFile("/etc/pam.d/su")
	if err != nil {
		return err
	}
	proposed = updatePAMsu(string(pamsu), enable)
	if proposed != string(pamsu) {
		if err := writeConfigFile("/etc/pam.d/su", proposed); err != nil {
			return err
		}
	}

	return nil
}

// Creates necessary OS Login directories if they don't exist.
func createOSLoginDirs() error {
	restorecon, restoreconerr := exec.LookPath("restorecon")

	for _, dir := range []string{"/var/google-sudoers.d", "/var/google-users.d"} {
		err := os.Mkdir(dir, 0750)
		if err != nil && !os.IsExist(err) {
			return err
		}
		if restoreconerr == nil {
			runCmd(exec.Command(restorecon, dir))
		}
	}
	return nil
}

func createOSLoginSudoersFile() error {
	osloginSudoers := "/etc/sudoers.d/google-oslogin"
	if runtime.GOOS == "freebsd" {
		osloginSudoers = "/usr/local" + osloginSudoers
	}
	sudoFile, err := os.OpenFile(osloginSudoers, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0440)
	if err != nil {
		if os.IsExist(err) {
			return nil
		}
		return err
	}
	fmt.Fprintf(sudoFile, "#includedir /var/google-sudoers.d\n")
	return sudoFile.Close()
}

// restartService tries to restart a service on linux-like systems. It attempts
// to find and use the following mechanisms in order:
// 1. The `systemctl` utility, if in a systemd environment.
// 2. The `service` command, if present.
// 3. A SysVinit script directly, if present.
// Missing mechanisms and missing or disabled services are ignored.
func restartService(servicename string) error {
	init, err := os.Readlink("/sbin/init")
	if err == nil && strings.Contains(init, "systemd") {
		if systemctl, err := exec.LookPath("systemctl"); err == nil {
			if err := runCmd(exec.Command(systemctl, "is-active", servicename+".service")); err == nil {
				return runCmd(exec.Command(systemctl, "restart", servicename+".service"))
			}
			return nil
		}
	}
	service, err := exec.LookPath("service")
	if err == nil {
		if err := runCmd(exec.Command(service, servicename, "status")); err == nil {
			return runCmd(exec.Command(service, servicename, "restart"))
		}
		return nil
	}
	initService := "/etc/init.d/" + servicename
	if _, err := os.Stat(initService); err == nil {
		if err := runCmd(exec.Command(initService, "status")); err == nil {
			return runCmd(exec.Command(initService, "restart"))
		}
		return nil
	}

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"sort"
	"strconv"
	"sync"
	"time"

	"github.com/lxc/lxd/lxd/db"
	"github.com/lxc/lxd/lxd/state"
	"github.com/lxc/lxd/shared"
	"github.com/lxc/lxd/shared/logger"

	log "github.com/lxc/lxd/shared/log15"
)

var containersCmd = Command{
	name: "containers",
	get:  containersGet,
	post: containersPost,
}

var containerCmd = Command{
	name:   "containers/{name}",
	get:    containerGet,
	put:    containerPut,
	delete: containerDelete,
	post:   containerPost,
	patch:  containerPatch,
}

var containerStateCmd = Command{
	name: "containers/{name}/state",
	get:  containerState,
	put:  containerStatePut,
}

var containerFileCmd = Command{
	name:   "containers/{name}/files",
	get:    containerFileHandler,
	post:   containerFileHandler,
	delete: containerFileHandler,
}

var containerSnapshotsCmd = Command{
	name: "containers/{name}/snapshots",
	get:  containerSnapshotsGet,
	post: containerSnapshotsPost,
}

var containerSnapshotCmd = Command{
	name:   "containers/{name}/snapshots/{snapshotName}",
	get:    snapshotHandler,
	post:   snapshotHandler,
	delete: snapshotHandler,
}

var containerConsoleCmd = Command{
	name:   "containers/{name}/console",
	get:    containerConsoleLogGet,
	post:   containerConsolePost,
	delete: containerConsoleLogDelete,
}

var containerExecCmd = Command{
	name: "containers/{name}/exec",
	post: containerExecPost,
}

var containerMetadataCmd = Command{
	name: "containers/{name}/metadata",
	get:  containerMetadataGet,
	put:  containerMetadataPut,
}

var containerMetadataTemplatesCmd = Command{
	name:   "containers/{name}/metadata/templates",
	get:    containerMetadataTemplatesGet,
	post:   containerMetadataTemplatesPostPut,
	put:    containerMetadataTemplatesPostPut,
	delete: containerMetadataTemplatesDelete,
}

type containerAutostartList []container

func (slice containerAutostartList) Len() int {
	return len(slice)
}

func (slice containerAutostartList) Less(i, j int) bool {
	iOrder := slice[i].ExpandedConfig()["boot.autostart.priority"]
	jOrder := slice[j].ExpandedConfig()["boot.autostart.priority"]

	if iOrder != jOrder {
		iOrderInt, _ := strconv.Atoi(iOrder)
		jOrderInt, _ := strconv.Atoi(jOrder)
		return iOrderInt > jOrderInt
	}

	return slice[i].Name() < slice[j].Name()
}

func (slice containerAutostartList) Swap(i, j int) {
	slice[i], slice[j] = slice[j], slice[i]
}

func containersRestart(s *state.State) error {
	// Get all the containers
	result, err := s.DB.ContainersList(db.CTypeRegular)
	if err != nil {
		return err
	}

	containers := []container{}

	for _, name := range result {
		c, err := containerLoadByName(s, name)
		if err != nil {
			return err
		}

		containers = append(containers, c)
	}

	sort.Sort(containerAutostartList(containers))

	// Restart the containers
	for _, c := range containers {
		config := c.ExpandedConfig()
		lastState := config["volatile.last_state.power"]

		autoStart := config["boot.autostart"]
		autoStartDelay := config["boot.autostart.delay"]

		if shared.IsTrue(autoStart) || (autoStart == "" && lastState == "RUNNING") {
			if c.IsRunning() {
				continue
			}

			c.Start(false)

			autoStartDelayInt, err := strconv.Atoi(autoStartDelay)
			if err == nil {
				time.Sleep(time.Duration(autoStartDelayInt) * time.Second)
			}
		}
	}

	return nil
}

type containerStopList []container

func (slice containerStopList) Len() int {
	return len(slice)
}

func (slice containerStopList) Less(i, j int) bool {
	iOrder := slice[i].ExpandedConfig()["boot.stop.priority"]
	jOrder := slice[j].ExpandedConfig()["boot.stop.priority"]

	if iOrder != jOrder {
		iOrderInt, _ := strconv.Atoi(iOrder)
		jOrderInt, _ := strconv.Atoi(jOrder)
		return iOrderInt > jOrderInt // check this line (prob <)
	}

	return slice[i].Name() < slice[j].Name()
}

func (slice containerStopList) Swap(i, j int) {
	slice[i], slice[j] = slice[j], slice[i]
}

func containersShutdown(s *state.State) error {
	var wg sync.WaitGroup

	// Get all the containers
	results, err := s.DB.ContainersList(db.CTypeRegular)
	if err != nil {
		return err
	}

	containers := []container{}

	for _, name := range results {
		c, err := containerLoadByName(s, name)
		if err != nil {
			return err
		}

		containers = append(containers, c)
	}

	sort.Sort(containerStopList(containers))

	// Reset all container states
	err = s.DB.ContainersResetState()
	if err != nil {
		return err
	}

	var lastPriority int = 0

	if len(containers) != 0 {
		lastPriority, _ = strconv.Atoi(containers[0].ExpandedConfig()["boot.stop.priority"])
	}

	for _, c := range containers {
		priority, _ := strconv.Atoi(c.ExpandedConfig()["boot.stop.priority"])

		// Enforce shutdown priority
		if priority != lastPriority {
			lastPriority = priority

			// Wait for containers with higher priority to finish
			wg.Wait()
		}

		// Record the current state
		lastState := c.State()

		// Stop the container
		if lastState != "BROKEN" && lastState != "STOPPED" {
			// Determinate how long to wait for the container to shutdown cleanly
			var timeoutSeconds int
			value, ok := c.ExpandedConfig()["boot.host_shutdown_timeout"]
			if ok {
				timeoutSeconds, _ = strconv.Atoi(value)
			} else {
				timeoutSeconds = 30
			}

			// Stop the container
			wg.Add(1)
			go func() {
				c.Shutdown(time.Second * time.Duration(timeoutSeconds))
				c.Stop(false)
				c.ConfigKeySet("volatile.last_state.power", lastState)

				wg.Done()
			}()
		} else {
			c.ConfigKeySet("volatile.last_state.power", lastState)
		}
	}
	wg.Wait()

	return nil
}

func containerDeleteSnapshots(s *state.State, cname string) error {
	logger.Debug("containerDeleteSnapshots",
		log.Ctx{"container": cname})

	results, err := s.DB.ContainerGetSnapshots(cname)
	if err != nil {
		return err
	}

	for _, sname := range results {
		sc, err := containerLoadByName(s, sname)
		if err != nil {
			logger.Error(
				"containerDeleteSnapshots: Failed to load the snapshotcontainer",
				log.Ctx{"container": cname, "snapshot": sname})

			continue
		}

		if err := sc.Delete(); err != nil {
			logger.Error(
				"containerDeleteSnapshots: Failed to delete a snapshotcontainer",
				log.Ctx{"container": cname, "snapshot": sname, "err": err})
		}
	}

	return nil
}
<commit_msg>lxd/containers: Log auto-start errors<commit_after>package main

import (
	"sort"
	"strconv"
	"sync"
	"time"

	"github.com/lxc/lxd/lxd/db"
	"github.com/lxc/lxd/lxd/state"
	"github.com/lxc/lxd/shared"
	"github.com/lxc/lxd/shared/logger"

	log "github.com/lxc/lxd/shared/log15"
)

var containersCmd = Command{
	name: "containers",
	get:  containersGet,
	post: containersPost,
}

var containerCmd = Command{
	name:   "containers/{name}",
	get:    containerGet,
	put:    containerPut,
	delete: containerDelete,
	post:   containerPost,
	patch:  containerPatch,
}

var containerStateCmd = Command{
	name: "containers/{name}/state",
	get:  containerState,
	put:  containerStatePut,
}

var containerFileCmd = Command{
	name:   "containers/{name}/files",
	get:    containerFileHandler,
	post:   containerFileHandler,
	delete: containerFileHandler,
}

var containerSnapshotsCmd = Command{
	name: "containers/{name}/snapshots",
	get:  containerSnapshotsGet,
	post: containerSnapshotsPost,
}

var containerSnapshotCmd = Command{
	name:   "containers/{name}/snapshots/{snapshotName}",
	get:    snapshotHandler,
	post:   snapshotHandler,
	delete: snapshotHandler,
}

var containerConsoleCmd = Command{
	name:   "containers/{name}/console",
	get:    containerConsoleLogGet,
	post:   containerConsolePost,
	delete: containerConsoleLogDelete,
}

var containerExecCmd = Command{
	name: "containers/{name}/exec",
	post: containerExecPost,
}

var containerMetadataCmd = Command{
	name: "containers/{name}/metadata",
	get:  containerMetadataGet,
	put:  containerMetadataPut,
}

var containerMetadataTemplatesCmd = Command{
	name:   "containers/{name}/metadata/templates",
	get:    containerMetadataTemplatesGet,
	post:   containerMetadataTemplatesPostPut,
	put:    containerMetadataTemplatesPostPut,
	delete: containerMetadataTemplatesDelete,
}

type containerAutostartList []container

func (slice containerAutostartList) Len() int {
	return len(slice)
}

func (slice containerAutostartList) Less(i, j int) bool {
	iOrder := slice[i].ExpandedConfig()["boot.autostart.priority"]
	jOrder := slice[j].ExpandedConfig()["boot.autostart.priority"]

	if iOrder != jOrder {
		iOrderInt, _ := strconv.Atoi(iOrder)
		jOrderInt, _ := strconv.Atoi(jOrder)
		return iOrderInt > jOrderInt
	}

	return slice[i].Name() < slice[j].Name()
}

func (slice containerAutostartList) Swap(i, j int) {
	slice[i], slice[j] = slice[j], slice[i]
}

func containersRestart(s *state.State) error {
	// Get all the containers
	result, err := s.DB.ContainersList(db.CTypeRegular)
	if err != nil {
		return err
	}

	containers := []container{}

	for _, name := range result {
		c, err := containerLoadByName(s, name)
		if err != nil {
			return err
		}

		containers = append(containers, c)
	}

	sort.Sort(containerAutostartList(containers))

	// Restart the containers
	for _, c := range containers {
		config := c.ExpandedConfig()
		lastState := config["volatile.last_state.power"]

		autoStart := config["boot.autostart"]
		autoStartDelay := config["boot.autostart.delay"]

		if shared.IsTrue(autoStart) || (autoStart == "" && lastState == "RUNNING") {
			if c.IsRunning() {
				continue
			}

			err = c.Start(false)
			if err != nil {
				logger.Errorf("Failed to start container '%s': %v", c.Name(), err)
			}

			autoStartDelayInt, err := strconv.Atoi(autoStartDelay)
			if err == nil {
				time.Sleep(time.Duration(autoStartDelayInt) * time.Second)
			}
		}
	}

	return nil
}

type containerStopList []container

func (slice containerStopList) Len() int {
	return len(slice)
}

func (slice containerStopList) Less(i, j int) bool {
	iOrder := slice[i].ExpandedConfig()["boot.stop.priority"]
	jOrder := slice[j].ExpandedConfig()["boot.stop.priority"]

	if iOrder != jOrder {
		iOrderInt, _ := strconv.Atoi(iOrder)
		jOrderInt, _ := strconv.Atoi(jOrder)
		return iOrderInt > jOrderInt // check this line (prob <)
	}

	return slice[i].Name() < slice[j].Name()
}

func (slice containerStopList) Swap(i, j int) {
	slice[i], slice[j] = slice[j], slice[i]
}

func containersShutdown(s *state.State) error {
	var wg sync.WaitGroup

	// Get all the containers
	results, err := s.DB.ContainersList(db.CTypeRegular)
	if err != nil {
		return err
	}

	containers := []container{}

	for _, name := range results {
		c, err := containerLoadByName(s, name)
		if err != nil {
			return err
		}

		containers = append(containers, c)
	}

	sort.Sort(containerStopList(containers))

	// Reset all container states
	err = s.DB.ContainersResetState()
	if err != nil {
		return err
	}

	var lastPriority int = 0

	if len(containers) != 0 {
		lastPriority, _ = strconv.Atoi(containers[0].ExpandedConfig()["boot.stop.priority"])
	}

	for _, c := range containers {
		priority, _ := strconv.Atoi(c.ExpandedConfig()["boot.stop.priority"])

		// Enforce shutdown priority
		if priority != lastPriority {
			lastPriority = priority

			// Wait for containers with higher priority to finish
			wg.Wait()
		}

		// Record the current state
		lastState := c.State()

		// Stop the container
		if lastState != "BROKEN" && lastState != "STOPPED" {
			// Determinate how long to wait for the container to shutdown cleanly
			var timeoutSeconds int
			value, ok := c.ExpandedConfig()["boot.host_shutdown_timeout"]
			if ok {
				timeoutSeconds, _ = strconv.Atoi(value)
			} else {
				timeoutSeconds = 30
			}

			// Stop the container
			wg.Add(1)
			go func() {
				c.Shutdown(time.Second * time.Duration(timeoutSeconds))
				c.Stop(false)
				c.ConfigKeySet("volatile.last_state.power", lastState)

				wg.Done()
			}()
		} else {
			c.ConfigKeySet("volatile.last_state.power", lastState)
		}
	}
	wg.Wait()

	return nil
}

func containerDeleteSnapshots(s *state.State, cname string) error {
	logger.Debug("containerDeleteSnapshots",
		log.Ctx{"container": cname})

	results, err := s.DB.ContainerGetSnapshots(cname)
	if err != nil {
		return err
	}

	for _, sname := range results {
		sc, err := containerLoadByName(s, sname)
		if err != nil {
			logger.Error(
				"containerDeleteSnapshots: Failed to load the snapshotcontainer",
				log.Ctx{"container": cname, "snapshot": sname})

			continue
		}

		if err := sc.Delete(); err != nil {
			logger.Error(
				"containerDeleteSnapshots: Failed to delete a snapshotcontainer",
				log.Ctx{"container": cname, "snapshot": sname, "err": err})
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/xml"
	"errors"
	"fmt"
	"net"
	"net/http"
	"strings"
)

type Server struct {
	Type           string `xml:"type,attr"`
	Hostname       string `xml:"hostname"`
	Port           uint16 `xml:"port"`
	SocketType     string `xml:"socketType"`
	Authentication string `xml:"authentication"`
	Username       string `xml:"username"`
}
type IncomingServer struct {
	XMLName struct{} `xml:"incomingServer"`

	Server
}
type OutgoingServer struct {
	XMLName struct{} `xml:"outgoingServer"`

	Server
}

type Provider struct {
	XMLName struct{} `xml:"emailProvider"`

	Id               string `xml:"id,attr"`
	Domain           string `xml:"domain"`
	DisplayName      string `xml:"displayName"`
	DisplayShortName string `xml:"displayShortName"`

	IncomingServers []IncomingServer
	OutgoingServers []OutgoingServer
}

type ClientConfig struct {
	XMLName struct{} `xml:"clientConfig"`

	Version   string `xml:"version,attr"`
	Providers []Provider
}

// Lookup the given service, protocol pair in the domain SRV records.
func lookup(service, proto, domain string) (string, uint16, error) {
	_, addresses, err := net.LookupSRV(service, proto, domain)

	if err != nil {
		return "", 0, err
	}
	if len(addresses) == 0 {
		return "", 0, errors.New("No SRV records available for the given domain")
	}

	return strings.Trim(addresses[0].Target, "."), addresses[0].Port, nil
}

// Generate an autoconfig XML document based on the information obtained from
// querying the domain SRV records.
func generate_xml(domain string) ([]byte, error) {
	// Incoming server.
	address_in, port_in, err := lookup("imaps", "tcp", domain)
	if err != nil {
		return nil, err
	}

	incoming := IncomingServer{}
	incoming.Type = "imap"
	incoming.Hostname = address_in
	incoming.Port = port_in
	incoming.SocketType = "SSL"
	incoming.Authentication = "password-cleartext"
	incoming.Username = "%EMAILLOCALPART%"

	// Outgoing server.
	address_out, port_out, err := lookup("submission", "tcp", domain)
	if err != nil {
		return nil, err
	}

	outgoing := OutgoingServer{}
	outgoing.Type = "smtp"
	outgoing.Hostname = address_out
	outgoing.Port = port_out
	outgoing.SocketType = "SSL"
	outgoing.Authentication = "password-cleartext"
	outgoing.Username = "%EMAILLOCALPART%"

	// Final data mangling.
	config := ClientConfig{
		Version: "1.1",
		Providers: []Provider{
			Provider{
				Id:               domain,
				Domain:           domain,
				DisplayName:      domain,
				DisplayShortName: domain,
				IncomingServers:  []IncomingServer{incoming},
				OutgoingServers:  []OutgoingServer{outgoing},
			},
		},
	}

	xmlconfig, err := xml.Marshal(&config)
	if err != nil {
		return nil, err
	}

	return xmlconfig, nil
}

func handler(w http.ResponseWriter, r *http.Request) {
	xmlconfig, err := generate_xml("marshland.ovh")
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	fmt.Fprint(w, xml.Header, string(xmlconfig))
}

func main() {
	http.HandleFunc("/", handler)
	http.ListenAndServe(":9090", nil)
}
<commit_msg>Created Domain struct to hold domain info.<commit_after>package main

import (
	"encoding/xml"
	"errors"
	"fmt"
	"net"
	"net/http"
	"strings"
)

type Server struct {
	Type           string `xml:"type,attr"`
	Hostname       string `xml:"hostname"`
	Port           uint16 `xml:"port"`
	SocketType     string `xml:"socketType"`
	Authentication string `xml:"authentication"`
	Username       string `xml:"username"`
}
type IncomingServer struct {
	XMLName struct{} `xml:"incomingServer"`

	Server
}
type OutgoingServer struct {
	XMLName struct{} `xml:"outgoingServer"`

	Server
}

type Provider struct {
	XMLName struct{} `xml:"emailProvider"`

	Id               string `xml:"id,attr"`
	Domain           string `xml:"domain"`
	DisplayName      string `xml:"displayName"`
	DisplayShortName string `xml:"displayShortName"`

	IncomingServers []IncomingServer
	OutgoingServers []OutgoingServer
}

type ClientConfig struct {
	XMLName struct{} `xml:"clientConfig"`

	Version   string `xml:"version,attr"`
	Providers []Provider
}

type Domain struct {
	domain string
	config ClientConfig
}

// Lookup the given service, protocol pair in the domain SRV records.
func (d *Domain) lookup(service, proto string) (string, uint16, error) {
	_, addresses, err := net.LookupSRV(service, proto, d.domain)

	if err != nil {
		return "", 0, err
	}
	if len(addresses) == 0 {
		return "", 0, errors.New("No SRV records available for the given domain")
	}

	return strings.Trim(addresses[0].Target, "."), addresses[0].Port, nil
}

// Generate an autoconfig XML document based on the information obtained from
// querying the domain SRV records.
func (d *Domain) generate_xml() ([]byte, error) {
	// Incoming server.
	address_in, port_in, err := d.lookup("imaps", "tcp")
	if err != nil {
		return nil, err
	}

	incoming := IncomingServer{}
	incoming.Type = "imap"
	incoming.Hostname = address_in
	incoming.Port = port_in
	incoming.SocketType = "SSL"
	incoming.Authentication = "password-cleartext"
	incoming.Username = "%EMAILLOCALPART%"

	// Outgoing server.
	address_out, port_out, err := d.lookup("submission", "tcp")
	if err != nil {
		return nil, err
	}

	outgoing := OutgoingServer{}
	outgoing.Type = "smtp"
	outgoing.Hostname = address_out
	outgoing.Port = port_out
	outgoing.SocketType = "SSL"
	outgoing.Authentication = "password-cleartext"
	outgoing.Username = "%EMAILLOCALPART%"

	// Final data mangling.
	d.config = ClientConfig{
		Version: "1.1",
		Providers: []Provider{
			Provider{
				Id:               d.domain,
				Domain:           d.domain,
				DisplayName:      d.domain,
				DisplayShortName: d.domain,
				IncomingServers:  []IncomingServer{incoming},
				OutgoingServers:  []OutgoingServer{outgoing},
			},
		},
	}

	xmlconfig, err := xml.Marshal(&d.config)
	if err != nil {
		return nil, err
	}

	return xmlconfig, nil
}

func handler(w http.ResponseWriter, r *http.Request) {
	domain := &Domain{"marshland.ovh", ClientConfig{}}
	xmlconfig, err := domain.generate_xml()
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	fmt.Fprint(w, xml.Header, string(xmlconfig))
}

func main() {
	http.HandleFunc("/", handler)
	http.ListenAndServe(":9090", nil)
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"os"
	"strings"

	"github.com/mkideal/cli"
)

const toolVersion = "0.0.4"

var verboseMode = false

type configuration struct {
	Sites []site
}

type site struct {
	Base      string
	Endpoints map[string]bool
}

type argT struct {
	cli.Helper
	Config  string `cli:"c,config" usage:"JSON config file, see config.json-template"`
	Verbose bool   `cli:"v,verbose" usage:"Verbose output"`
	Version bool   `cli:"version" usage:"Check version"`
}

func getMaxWidth(sites []site) (width int) {
	var URL string
	for _, site := range sites {
		for endpoint := range site.Endpoints {
			URL = fmt.Sprintf("%s/%s", site.Base, endpoint)
			if len(URL) > width {
				width = len(URL)
			}
		}
	}
	return width
}

func checkSites(sites []site) error {
	maxWidth := getMaxWidth(sites)
	if !verboseMode {
		fmt.Printf("%*s | %*s | %*s | HTTP Status\n%s-+-%s-+-%s-+-%s\n", maxWidth, "URL", 10, "Basic Auth",
			16, "Success", strings.Repeat("-", maxWidth), strings.Repeat("-", 10),
			strings.Repeat("-", 16), strings.Repeat("-", 80-maxWidth-2))
	}
	for _, site := range sites {
		if verboseMode {
			log.Printf("Checking site %s\n", site.Base)
		}
		// channel for synchronizing 'done state', buffer the amount of endpoints
		done := make(chan bool, len(site.Endpoints))
		for endpoint, baShouldBe := range site.Endpoints {
			go checkEndpoint(done, &site, endpoint, baShouldBe, maxWidth)
		}
		// Drain the channel and wait for all goroutines to complete
		for i := 0; i < len(site.Endpoints); i++ {
			<-done // wait for one task to complete
		}
	}
	return nil
}

func checkSuccess(response *http.Response, baShouldBe bool) (success bool, baEnabled bool) {
	if response.StatusCode == 401 {
		baEnabled = true
	}
	return baEnabled == baShouldBe, response.StatusCode == 401
}

func checkEndpoint(done chan bool, site *site, endpoint string, baShouldBe bool, maxWidth int) error {
	response, err := getEndpoint(site.Base, endpoint, baShouldBe)
	if err != nil {
		return err
	}
	success, baEnabled := checkSuccess(response, baShouldBe)
	var message string
	var logMessage string

	if success {
		message = fmt.Sprintf("%*s | %*s | %*t | %s\n", maxWidth, fmt.Sprintf("%s/%s", site.Base, endpoint), 10, "yes", 16, success, response.Status)
		logMessage = fmt.Sprintf("OK: %s/%s correct. Basic Auth Enabled: %t Should be: %t\n", site.Base, endpoint, baEnabled, baShouldBe)
	} else {
		if response.StatusCode > 401 {
			message = fmt.Sprintf("%*s | %*s | %*t | %s\n", maxWidth, fmt.Sprintf("%s/%s", site.Base, endpoint), 10, "unknown", 16, success, response.Status)
			logMessage = fmt.Sprintf("ERROR: %s/%s unknown. %s\n", site.Base, endpoint, response.Status)
		} else {
			message = fmt.Sprintf("%*s | %*s | %*t | %s\n", maxWidth, fmt.Sprintf("%s/%s", site.Base, endpoint), 10, "no", 16, success, response.Status)
			logMessage = fmt.Sprintf("ERROR: %s/%s incorrect. Basic Auth Enabled: %t Should be: %t\n", site.Base, endpoint, baEnabled, baShouldBe)
		}
	}

	if verboseMode {
		log.Printf(logMessage)
	} else {
		fmt.Printf(message)
	}
	done <- true
	return nil
}

func getEndpoint(site string, endpoint string, baShouldBe bool) (*http.Response, error) {
	if verboseMode {
		log.Printf("Checking endpoint %v, Basic Auth should be %v\n", endpoint, baShouldBe)
	}
	URL := fmt.Sprintf("%s/%s", site, endpoint)
	resp, err := http.Get(URL)
	return resp, err
}

func main() {
	cli.Run(&argT{}, func(ctx *cli.Context) error {
		argv := ctx.Argv().(*argT)
		if argv.Version {
			ctx.String("ba_checker v%s\n", toolVersion)
			return nil
		}
		if argv.Config == "" {
			return fmt.Errorf("--config <config.json> is required.\n")
		}
		if argv.Verbose {
			verboseMode = true
		}
		if _, err := os.Stat(argv.Config); os.IsNotExist(err) {
			return fmt.Errorf("Error: %s does not exist", argv.Config)
		}
		file, _ := os.Open(argv.Config)
		decoder := json.NewDecoder(file)
		config := configuration{}
		err := decoder.Decode(&config)
		if err != nil {
			fmt.Println("error:", err)
		}
		err = checkSites(config.Sites)
		if err != nil {
			return err
		}
		return nil
	})
}
<commit_msg>Symbolic version bump<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"os"
	"strings"

	"github.com/mkideal/cli"
)

const toolVersion = "0.1"

var verboseMode = false

type configuration struct {
	Sites []site
}

type site struct {
	Base      string
	Endpoints map[string]bool
}

type argT struct {
	cli.Helper
	Config  string `cli:"c,config" usage:"JSON config file, see config.json-template"`
	Verbose bool   `cli:"v,verbose" usage:"Verbose output"`
	Version bool   `cli:"version" usage:"Check version"`
}

func getMaxWidth(sites []site) (width int) {
	var URL string
	for _, site := range sites {
		for endpoint := range site.Endpoints {
			URL = fmt.Sprintf("%s/%s", site.Base, endpoint)
			if len(URL) > width {
				width = len(URL)
			}
		}
	}
	return width
}

func checkSites(sites []site) error {
	maxWidth := getMaxWidth(sites)
	if !verboseMode {
		fmt.Printf("%*s | %*s | %*s | HTTP Status\n%s-+-%s-+-%s-+-%s\n", maxWidth, "URL", 10, "Basic Auth",
			16, "Success", strings.Repeat("-", maxWidth), strings.Repeat("-", 10),
			strings.Repeat("-", 16), strings.Repeat("-", 80-maxWidth-2))
	}
	for _, site := range sites {
		if verboseMode {
			log.Printf("Checking site %s\n", site.Base)
		}
		// channel for synchronizing 'done state', buffer the amount of endpoints
		done := make(chan bool, len(site.Endpoints))
		for endpoint, baShouldBe := range site.Endpoints {
			go checkEndpoint(done, &site, endpoint, baShouldBe, maxWidth)
		}
		// Drain the channel and wait for all goroutines to complete
		for i := 0; i < len(site.Endpoints); i++ {
			<-done // wait for one task to complete
		}
	}
	return nil
}

func checkSuccess(response *http.Response, baShouldBe bool) (success bool, baEnabled bool) {
	if response.StatusCode == 401 {
		baEnabled = true
	}
	return baEnabled == baShouldBe, response.StatusCode == 401
}

func checkEndpoint(done chan bool, site *site, endpoint string, baShouldBe bool, maxWidth int) error {
	response, err := getEndpoint(site.Base, endpoint, baShouldBe)
	if err != nil {
		return err
	}
	success, baEnabled := checkSuccess(response, baShouldBe)
	var message string
	var logMessage string

	if success {
		message = fmt.Sprintf("%*s | %*s | %*t | %s\n", maxWidth, fmt.Sprintf("%s/%s", site.Base, endpoint), 10, "yes", 16, success, response.Status)
		logMessage = fmt.Sprintf("OK: %s/%s correct. Basic Auth Enabled: %t Should be: %t\n", site.Base, endpoint, baEnabled, baShouldBe)
	} else {
		if response.StatusCode > 401 {
			message = fmt.Sprintf("%*s | %*s | %*t | %s\n", maxWidth, fmt.Sprintf("%s/%s", site.Base, endpoint), 10, "unknown", 16, success, response.Status)
			logMessage = fmt.Sprintf("ERROR: %s/%s unknown. %s\n", site.Base, endpoint, response.Status)
		} else {
			message = fmt.Sprintf("%*s | %*s | %*t | %s\n", maxWidth, fmt.Sprintf("%s/%s", site.Base, endpoint), 10, "no", 16, success, response.Status)
			logMessage = fmt.Sprintf("ERROR: %s/%s incorrect. Basic Auth Enabled: %t Should be: %t\n", site.Base, endpoint, baEnabled, baShouldBe)
		}
	}

	if verboseMode {
		log.Printf(logMessage)
	} else {
		fmt.Printf(message)
	}
	done <- true
	return nil
}

func getEndpoint(site string, endpoint string, baShouldBe bool) (*http.Response, error) {
	if verboseMode {
		log.Printf("Checking endpoint %v, Basic Auth should be %v\n", endpoint, baShouldBe)
	}
	URL := fmt.Sprintf("%s/%s", site, endpoint)
	resp, err := http.Get(URL)
	return resp, err
}

func main() {
	cli.Run(&argT{}, func(ctx *cli.Context) error {
		argv := ctx.Argv().(*argT)
		if argv.Version {
			ctx.String("ba_checker v%s\n", toolVersion)
			return nil
		}
		if argv.Config == "" {
			return fmt.Errorf("--config <config.json> is required.\n")
		}
		if argv.Verbose {
			verboseMode = true
		}
		if _, err := os.Stat(argv.Config); os.IsNotExist(err) {
			return fmt.Errorf("Error: %s does not exist", argv.Config)
		}
		file, _ := os.Open(argv.Config)
		decoder := json.NewDecoder(file)
		config := configuration{}
		err := decoder.Decode(&config)
		if err != nil {
			fmt.Println("error:", err)
		}
		err = checkSites(config.Sites)
		if err != nil {
			return err
		}
		return nil
	})
}
<|endoftext|>


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

package zaptest

import (
	"go.uber.org/zap"
	"go.uber.org/zap/zapcore"
)

// LoggerOption configures the test logger built by NewLogger.
type LoggerOption interface {
	applyLoggerOption(*loggerOptions)
}

type loggerOptions struct {
	Level zapcore.LevelEnabler
}

type loggerOptionFunc func(*loggerOptions)

func (f loggerOptionFunc) applyLoggerOption(opts *loggerOptions) {
	f(opts)
}

// Level controls which messages are logged by a test Logger built by
// NewLogger.
func Level(enab zapcore.LevelEnabler) LoggerOption {
	return loggerOptionFunc(func(opts *loggerOptions) {
		opts.Level = enab
	})
}

// NewLogger builds a new Logger that logs all messages to the given
// testing.TB.
//
//   logger := zaptest.NewLogger(t, zaptest.Level(zap.WarnLevel))
//
// Use this with a *testing.T or *testing.B to get logs which get printed only
// if a test fails or if you ran go test -v.
func NewLogger(t TestingT, opts ...LoggerOption) *zap.Logger {
	cfg := loggerOptions{
		Level: zapcore.DebugLevel,
	}
	for _, o := range opts {
		o.applyLoggerOption(&cfg)
	}

	return zap.New(
		zapcore.NewCore(
			zapcore.NewConsoleEncoder(zap.NewDevelopmentEncoderConfig()),
			testingWriter{t},
			cfg.Level,
		),
	)
}

// testingWriter is a WriteSyncer that writes to the given testing.TB.
type testingWriter struct{ t TestingT }

func (w testingWriter) Write(p []byte) (n int, err error) {
	s := string(p)

	// Strip trailing newline because t.Log always adds one.
	if s[len(s)-1] == '\n' {
		s = s[:len(s)-1]
	}

	// Note: t.Log is safe for concurrent use.
	w.t.Logf("%s", s)
	return len(p), nil
}

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

package zaptest

import (
	"bytes"

	"go.uber.org/zap"
	"go.uber.org/zap/zapcore"
)

// LoggerOption configures the test logger built by NewLogger.
type LoggerOption interface {
	applyLoggerOption(*loggerOptions)
}

type loggerOptions struct {
	Level zapcore.LevelEnabler
}

type loggerOptionFunc func(*loggerOptions)

func (f loggerOptionFunc) applyLoggerOption(opts *loggerOptions) {
	f(opts)
}

// Level controls which messages are logged by a test Logger built by
// NewLogger.
func Level(enab zapcore.LevelEnabler) LoggerOption {
	return loggerOptionFunc(func(opts *loggerOptions) {
		opts.Level = enab
	})
}

// NewLogger builds a new Logger that logs all messages to the given
// testing.TB.
//
//   logger := zaptest.NewLogger(t, zaptest.Level(zap.WarnLevel))
//
// Use this with a *testing.T or *testing.B to get logs which get printed only
// if a test fails or if you ran go test -v.
func NewLogger(t TestingT, opts ...LoggerOption) *zap.Logger {
	cfg := loggerOptions{
		Level: zapcore.DebugLevel,
	}
	for _, o := range opts {
		o.applyLoggerOption(&cfg)
	}

	return zap.New(
		zapcore.NewCore(
			zapcore.NewConsoleEncoder(zap.NewDevelopmentEncoderConfig()),
			testingWriter{t},
			cfg.Level,
		),
	)
}

// testingWriter is a WriteSyncer that writes to the given testing.TB.
type testingWriter struct{ t TestingT }

func (w testingWriter) Write(p []byte) (n int, err error) {
	// Strip trailing newline because t.Log always adds one.
	p = bytes.TrimRight(p, "\n")

	// Note: t.Log is safe for concurrent use.
	w.t.Logf("%s", p)
	return len(p), nil
}

func (w testingWriter) Sync() error {
	return nil
}
<|endoftext|>


<commit_before>package profile

import (
	"bufio"
	"bytes"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"testing"
)

type checkFn func(t *testing.T, stdout, stderr []byte, err error)

var profileTests = []struct {
	name   string
	code   string
	checks []checkFn
}{{
	name: "default profile (cpu)",
	code: `
package main

import "github.com/pkg/profile"

func main() {
	defer profile.Start().Stop()
}
`,
	checks: []checkFn{
		NoStdout,
		Stderr("profile: cpu profiling enabled"),
		NoErr,
	},
}, {
	name: "memory profile",
	code: `
package main

import "github.com/pkg/profile"

func main() {
	defer profile.Start(profile.MemProfile).Stop()
}
`,
	checks: []checkFn{
		NoStdout,
		Stderr("profile: memory profiling enabled"),
		NoErr,
	},
}, {
	name: "memory profile (rate 2048)",
	code: `
package main

import "github.com/pkg/profile"

func main() {
	defer profile.Start(profile.MemProfileRate(2048)).Stop()
}
`,
	checks: []checkFn{
		NoStdout,
		Stderr("profile: memory profiling enabled (rate 2048)"),
		NoErr,
	},
}, {
	name: "double start",
	code: `
package main

import "github.com/pkg/profile"

func main() {
	defer profile.Start().Stop()
	defer profile.Start().Stop()
}
`,
	checks: []checkFn{
		NoStdout,
		Stderr("cpu profiling enabled", "profile: Start() already called"),
		Err,
	},
}, {
	name: "block profile",
	code: `
package main

import "github.com/pkg/profile"

func main() {
	defer profile.Start(profile.BlockProfile).Stop()
}
`,
	checks: []checkFn{
		NoStdout,
		Stderr("profile: block profiling enabled"),
		NoErr,
	},
}, {
	name: "profile path",
	code: `
package main

import "github.com/pkg/profile"

func main() {
	defer profile.Start(profile.ProfilePath(".")).Stop()
}
`,
	checks: []checkFn{
		NoStdout,
		Stderr("profile: cpu profiling enabled, cpu.pprof"),
		NoErr,
	},
}, {
	name: "profile path error",
	code: `
package main

import "github.com/pkg/profile"

func main() {
		defer profile.Start(profile.ProfilePath("README.md")).Stop()
}
`,
	checks: []checkFn{
		NoStdout,
		Stderr("could not create initial output"),
		Err,
	},
}, {
	name: "profile quiet",
	code: `
package main

import "github.com/pkg/profile"

func main() {
        defer profile.Start(profile.Quiet).Stop()
}
`,
	checks: []checkFn{NoStdout, NoStderr, NoErr},
}}

func TestProfile(t *testing.T) {
	for _, tt := range profileTests {
		t.Log(tt.name)
		stdout, stderr, err := runTest(t, tt.code)
		for _, f := range tt.checks {
			f(t, stdout, stderr, err)
		}
	}
}

// NoStdout checks that stdout was blank.
func NoStdout(t *testing.T, stdout, _ []byte, _ error) {
	if len := len(stdout); len > 0 {
		t.Errorf("stdout: wanted 0 bytes, got %d", len)
	}
}

// Stderr verifies that the given lines match the output from stderr
func Stderr(lines ...string) checkFn {
	return func(t *testing.T, stdout, stderr []byte, _ error) {
		buf := bytes.NewBuffer(stderr)
		if !validateOutput(buf, lines) {
			t.Errorf("stderr: wanted '%s', got '%s'", lines, stderr)
		}
	}
}

// NoStderr checks that stderr was blank.
func NoStderr(t *testing.T, _, stderr []byte, _ error) {
	if len := len(stderr); len > 0 {
		t.Errorf("stderr: wanted 0 bytes, got %d", len)
	}
}

// Err checks that there was an error returned
func Err(t *testing.T, _, _ []byte, err error) {
	if err == nil {
		t.Errorf("expected error")
	}
}

// NoErr checks that err was nil
func NoErr(t *testing.T, _, _ []byte, err error) {
	if err != nil {
		t.Errorf("error: expected nil, got %v", err)
	}
}

// validatedOutput validates the given slice of lines against a scanned buffer
func validateOutput(buf *bytes.Buffer, lines []string) bool {
	s := bufio.NewScanner(buf)
	for _, l := range lines {
		if !s.Scan() || !strings.Contains(s.Text(), l) {
			return false
		}
	}
	return true
}

// runTest executes the go program supplied and returns the contents of stdout,
// stderr, and an error which may contain status information about the result
// of the program.
func runTest(t *testing.T, code string) ([]byte, []byte, error) {
	chk := func(err error) {
		if err != nil {
			t.Fatal(err)
		}
	}
	gopath, err := ioutil.TempDir("", "profile-gopath")
	chk(err)
	defer os.RemoveAll(gopath)

	srcdir := filepath.Join(gopath, "src")
	err = os.Mkdir(srcdir, 0755)
	chk(err)
	src := filepath.Join(srcdir, "main.go")
	err = ioutil.WriteFile(src, []byte(code), 0644)
	chk(err)

	cmd := exec.Command("go", "run", src)

	var stdout, stderr bytes.Buffer
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr
	err = cmd.Run()
	return stdout.Bytes(), stderr.Bytes(), err
}
<commit_msg>Addressed review comments.<commit_after>package profile

import (
	"bufio"
	"bytes"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"testing"
)

type checkFn func(t *testing.T, stdout, stderr []byte, err error)

var profileTests = []struct {
	name   string
	code   string
	checks []checkFn
}{{
	name: "default profile (cpu)",
	code: `
package main

import "github.com/pkg/profile"

func main() {
	defer profile.Start().Stop()
}
`,
	checks: []checkFn{
		NoStdout,
		Stderr("profile: cpu profiling enabled"),
		NoErr,
	},
}, {
	name: "memory profile",
	code: `
package main

import "github.com/pkg/profile"

func main() {
	defer profile.Start(profile.MemProfile).Stop()
}
`,
	checks: []checkFn{
		NoStdout,
		Stderr("profile: memory profiling enabled"),
		NoErr,
	},
}, {
	name: "memory profile (rate 2048)",
	code: `
package main

import "github.com/pkg/profile"

func main() {
	defer profile.Start(profile.MemProfileRate(2048)).Stop()
}
`,
	checks: []checkFn{
		NoStdout,
		Stderr("profile: memory profiling enabled (rate 2048)"),
		NoErr,
	},
}, {
	name: "double start",
	code: `
package main

import "github.com/pkg/profile"

func main() {
	defer profile.Start().Stop()
	defer profile.Start().Stop()
}
`,
	checks: []checkFn{
		NoStdout,
		Stderr("cpu profiling enabled", "profile: Start() already called"),
		Err,
	},
}, {
	name: "block profile",
	code: `
package main

import "github.com/pkg/profile"

func main() {
	defer profile.Start(profile.BlockProfile).Stop()
}
`,
	checks: []checkFn{
		NoStdout,
		Stderr("profile: block profiling enabled"),
		NoErr,
	},
}, {
	name: "profile path",
	code: `
package main

import "github.com/pkg/profile"

func main() {
	defer profile.Start(profile.ProfilePath(".")).Stop()
}
`,
	checks: []checkFn{
		NoStdout,
		Stderr("profile: cpu profiling enabled, cpu.pprof"),
		NoErr,
	},
}, {
	name: "profile path error",
	code: `
package main

import "github.com/pkg/profile"

func main() {
		defer profile.Start(profile.ProfilePath("README.md")).Stop()
}
`,
	checks: []checkFn{
		NoStdout,
		Stderr("could not create initial output"),
		Err,
	},
}, {
	name: "profile quiet",
	code: `
package main

import "github.com/pkg/profile"

func main() {
        defer profile.Start(profile.Quiet).Stop()
}
`,
	checks: []checkFn{NoStdout, NoStderr, NoErr},
}}

func TestProfile(t *testing.T) {
	for _, tt := range profileTests {
		t.Log(tt.name)
		stdout, stderr, err := runTest(t, tt.code)
		for _, f := range tt.checks {
			f(t, stdout, stderr, err)
		}
	}
}

// NoStdout checks that stdout was blank.
func NoStdout(t *testing.T, stdout, _ []byte, _ error) {
	if len := len(stdout); len > 0 {
		t.Errorf("stdout: wanted 0 bytes, got %d", len)
	}
}

// Stderr verifies that the given lines match the output from stderr
func Stderr(lines ...string) checkFn {
	return func(t *testing.T, stdout, stderr []byte, _ error) {
		r := bytes.NewReader(stderr)
		if !validateOutput(r, lines) {
			t.Errorf("stderr: wanted '%s', got '%s'", lines, stderr)
		}
	}
}

// NoStderr checks that stderr was blank.
func NoStderr(t *testing.T, _, stderr []byte, _ error) {
	if len := len(stderr); len > 0 {
		t.Errorf("stderr: wanted 0 bytes, got %d", len)
	}
}

// Err checks that there was an error returned
func Err(t *testing.T, _, _ []byte, err error) {
	if err == nil {
		t.Errorf("expected error")
	}
}

// NoErr checks that err was nil
func NoErr(t *testing.T, _, _ []byte, err error) {
	if err != nil {
		t.Errorf("error: expected nil, got %v", err)
	}
}

// validatedOutput validates the given slice of lines against data from the given reader.
func validateOutput(r *bytes.Reader, want []string) bool {
	s := bufio.NewScanner(r)
	for _, line := range want {
		if !s.Scan() || !strings.Contains(s.Text(), line) {
			return false
		}
	}
	return true
}

// runTest executes the go program supplied and returns the contents of stdout,
// stderr, and an error which may contain status information about the result
// of the program.
func runTest(t *testing.T, code string) ([]byte, []byte, error) {
	chk := func(err error) {
		if err != nil {
			t.Fatal(err)
		}
	}
	gopath, err := ioutil.TempDir("", "profile-gopath")
	chk(err)
	defer os.RemoveAll(gopath)

	srcdir := filepath.Join(gopath, "src")
	err = os.Mkdir(srcdir, 0755)
	chk(err)
	src := filepath.Join(srcdir, "main.go")
	err = ioutil.WriteFile(src, []byte(code), 0644)
	chk(err)

	cmd := exec.Command("go", "run", src)

	var stdout, stderr bytes.Buffer
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr
	err = cmd.Run()
	return stdout.Bytes(), stderr.Bytes(), err
}
<|endoftext|>


<commit_before>/*

In a number of cases you have a property that can only have a handful of
possible values. You want to verify that the value is always one of those
legal values, and make sure that you can compare it to a known constant so you
can make sure you don't have a typo at compile time instead of run time. It's
also nice to have them have an order in many cases, and to be serialized with
the string value so it's easier to read.

Enums are useful for this case. An EnumSet contains multiple enums, and you
can create an EnumValue which can be used as a property on a PropertyReader
object.

The idiomatic way to create an enum is the following:
	const (
		//The first enum can start at 0
		ColorRed = iota
		ColorBlue
		ColorGreen
	)

	const (
		//The second enum should start at 1 plus the last item in the
		//previous, because all int vals in an EnumSet must be unique.
		CardSpade = ColorGreen + 1 + iota
		CardHeart
		CardDiamond
		CardClub
	)

	var Enums = enum.NewSet()

	var ColorEnum = Enums.MustAdd("Color", map[int]string{
		ColorRed: "Red",
		ColorBlue: "Blue",
		ColorGreen: "Green",
	})

	var CardEnum = Enums.MustAdd("Card", map[int]string{
		CardSpade: "Spade",
		CardHeart: "Heart",
		CardDiamond: "Diamond",
		CardClub: "Club",
	})

	//...

	func (g *GameDelegate) EmptyGameState() boardgame.SubState {
		return &gameState{
			MyIntProp: 0,
			MyColorEnumProp: ColorEnum.NewEnumValue(),
		}
	}

	//...

	func NewManager() *boardgame.GameManager {
		//...

		//NewComponentChest will call Finish() on our Enums
		chest := boardgame.NewComponentChest(Enums)

		//...
	}

*/
package enum

import (
	"encoding/json"
	"errors"
	"math"
	"strconv"
)

//EnumSet is a set of enums where each Enum's values are unique. Normally you
//will create one in your package, add enums to it during initalization, and
//then use it for all managers you create.
type Set struct {
	finished bool
	enums    map[string]*Enum
	//A map of which int goes to which Enum
	values map[int]*Enum
}

//Enum is a named set of values within a set. Get a new one with
//enumSet.Add().
type Enum struct {
	name         string
	values       map[int]string
	defaultValue int
}

//variable is the underlying type we'll return for both Value and Constant.
type variable struct {
	enum *Enum
	val  int
}

type Const interface {
	Enum() *Enum
	Value() int
	String() string
	Copy() Const
}

//An EnumValue is an instantiation of a value that must be set to a value in
//the given enum. You retrieve one from enum.NewEnumValue().
type Var interface {
	Const
	SetValue(int) error
	CopyVar() Var
}

//NewSet returns a new Set. Generally you'll call this once in a
//package and create the set during initalization.
func NewSet() *Set {
	return &Set{
		false,
		make(map[string]*Enum),
		make(map[int]*Enum),
	}
}

//MustCombineSets wraps CombineEnumSets, but instead of erroring will
//panic. Useful for package-level declarations outside of init().
func MustCombineSets(sets ...*Set) *Set {
	result, err := CombineSets(sets...)
	if err != nil {
		panic("Couldn't combine sets: " + err.Error())
	}
	return result
}

//CombineSets returns a new EnumSet that contains all of the EnumSets
//combined into one. The individual enums will literally be the same as the
//enums from the provided sets, so enum equality will work. Generally the sets
//have to know about each other, otherwise they are liable to overlap, which
//will error.
func CombineSets(sets ...*Set) (*Set, error) {
	result := NewSet()
	for i, set := range sets {
		for _, enumName := range set.EnumNames() {
			enum := set.Enum(enumName)
			if err := result.addEnum(enumName, enum); err != nil {
				return nil, errors.New("Couldn't add the " + strconv.Itoa(i) + " enumset because " + enumName + " had error: " + err.Error())
			}
		}
	}
	return result, nil
}

//Finish finalizes an EnumSet so that no more enums may be added. After this
//is called it is safe to use this in a multi-threaded environment. Repeated
//calls do nothing. ComponenChest automatically calls Finish() on the set you
//pass it.
func (e *Set) Finish() {
	e.finished = true
}

//EnumNames returns a list of all names in the Enum.
func (e *Set) EnumNames() []string {
	var result []string
	for key, _ := range e.enums {
		result = append(result, key)
	}
	return result
}

//Returns the Enum with the given name. In general you keep a reference to the
//enum yourself, but this is useful for programatically enumerating the enums.
func (e *Set) Enum(name string) *Enum {
	return e.enums[name]
}

//Membership returns the enum that the given val is a member of.
func (e *Set) Membership(val int) *Enum {
	return e.values[val]
}

//MustAdd is like Add, but instead of an error it will panic if the enum
//cannot be added. This is useful for defining your enums at the package level
//outside of an init().
func (e *Set) MustAdd(enumName string, values map[int]string) *Enum {
	result, err := e.Add(enumName, values)

	if err != nil {
		panic("Couldn't add to enumset: " + err.Error())
	}

	return result
}

/*
Add ads an enum with the given name and values to the enum manager. Will error
if that name has already been added, or any of the int values has been used
for any other enum item already. This means that enums must be unique within a
manager. Check out the package doc for the idiomatic way to initalize enums.
*/
func (e *Set) Add(enumName string, values map[int]string) (*Enum, error) {

	if len(values) == 0 {
		return nil, errors.New("No values provided")
	}

	enum := &Enum{
		enumName,
		make(map[int]string),
		math.MaxInt64,
	}

	seenValues := make(map[string]bool)

	for v, s := range values {

		if seenValues[s] {
			return nil, errors.New("String " + s + " was not unique within enum " + enumName)
		}

		seenValues[s] = true

		enum.values[v] = s

		if v < enum.defaultValue {
			enum.defaultValue = v
		}

	}
	if err := e.addEnum(enumName, enum); err != nil {
		return nil, err
	}
	return enum, nil
}

func (e *Set) addEnum(enumName string, enum *Enum) error {

	if e.finished {
		return errors.New("The set has been finished so no more enums can be added")
	}

	if _, ok := e.enums[enumName]; ok {
		return errors.New("That enum name has already been provided")
	}

	for v, _ := range enum.values {
		if _, ok := e.values[v]; ok {
			//Already registered
			return errors.New("Value " + strconv.Itoa(v) + " was registered twice")
		}

		e.values[v] = enum
	}

	e.enums[enumName] = enum

	return nil
}

//DefaultValue returns the default value for this enum (the lowest valid value
//in it).
func (e *Enum) DefaultValue() int {
	return e.defaultValue
}

//Valid returns whether the given value is a valid member of this enum.
func (e *Enum) Valid(val int) bool {
	_, ok := e.values[val]
	return ok
}

//String returns the string value associated with the given value.
func (e *Enum) String(val int) string {
	return e.values[val]
}

//Name returns the name of this enum; if set is the set this enum is part of,
//set.Enum(enum.Name()) == enum will be true.
func (e *Enum) Name() string {
	return e.name
}

//ValueFromString returns the enum value that corresponds to the given string,
//or -1 if no value has that string.
func (e *Enum) ValueFromString(in string) int {
	for v, str := range e.values {
		if str == in {
			return v
		}
	}
	return -1
}

//Copy returns a copy of the Value, that is equivalent, but will not be
//locked.
func (e *variable) Copy() Const {
	return &variable{
		e.enum,
		e.val,
	}
}

func (e *variable) CopyVar() Var {
	return &variable{
		e.enum,
		e.val,
	}
}

//NewEnumValue returns a new EnumValue associated with this enum, set to the
//Enum's DefaultValue to start.
func (e *Enum) NewVar() Var {
	return &variable{
		e,
		e.DefaultValue(),
	}
}

//MustNewConst is like NewConst, but if it would have errored it panics
//instead. It's convenient for initial set up where the whole app should fail
//to startup if it can't be configured anyway, and dealing with errors would
//be a lot of boilerplate.
func (e *Enum) MustNewConst(val int) Const {
	result, err := e.NewConst(val)
	if err != nil {
		panic("Couldn't create constant: " + err.Error())
	}
	return result
}

//NewConstant returns an enum.Constant that is permanently set to the provided
//val. If that value is not valid for this enum, it will error.
func (e *Enum) NewConst(val int) (Const, error) {
	variable := e.NewVar()
	if err := variable.SetValue(val); err != nil {
		return nil, err
	}
	return variable, nil
}

//The enum marshals as the string value of the enum so it's more readable.
func (e *variable) MarshalJSON() ([]byte, error) {
	return json.Marshal(e.String())
}

//UnmarshalJSON expects the blob to be the string value. Will error if that
//doesn't correspond to a valid value for this enum.
func (e *variable) UnmarshalJSON(blob []byte) error {
	var str string
	if err := json.Unmarshal(blob, &str); err != nil {
		return err
	}
	val := e.enum.ValueFromString(str)
	if val == -1 {
		return errors.New("That string value had no enum in the value")
	}
	return e.SetValue(val)
}

func (e *variable) Enum() *Enum {
	return e.enum
}

func (e *variable) Value() int {
	return e.val
}

func (e *variable) String() string {
	return e.enum.String(e.val)
}

//SetValue changes the value. Returns true if successful. Will fail if the
//value is locked or the val you want to set is not a valid number for the
//enum this value is associated with.
func (e *variable) SetValue(val int) error {
	if !e.enum.Valid(val) {
		return errors.New("That value is not valid for this enum")
	}
	e.val = val
	return nil
}
<commit_msg>Added a bit more package doc for enum.Const and enum.Var. Part of #457.<commit_after>/*

In a number of cases you have a property that can only have a handful of
possible values. You want to verify that the value is always one of those
legal values, and make sure that you can compare it to a known constant so you
can make sure you don't have a typo at compile time instead of run time. It's
also nice to have them have an order in many cases, and to be serialized with
the string value so it's easier to read.

Enums are useful for this case. An EnumSet contains multiple enums, and you
can create an EnumValue which can be used as a property on a PropertyReader
object.

The idiomatic way to create an enum is the following:
	const (
		//The first enum can start at 0
		ColorRed = iota
		ColorBlue
		ColorGreen
	)

	const (
		//The second enum should start at 1 plus the last item in the
		//previous, because all int vals in an EnumSet must be unique.
		CardSpade = ColorGreen + 1 + iota
		CardHeart
		CardDiamond
		CardClub
	)

	var Enums = enum.NewSet()

	var ColorEnum = Enums.MustAdd("Color", map[int]string{
		ColorRed: "Red",
		ColorBlue: "Blue",
		ColorGreen: "Green",
	})

	var CardEnum = Enums.MustAdd("Card", map[int]string{
		CardSpade: "Spade",
		CardHeart: "Heart",
		CardDiamond: "Diamond",
		CardClub: "Club",
	})

	//...

	func (g *GameDelegate) EmptyGameState() boardgame.SubState {
		return &gameState{
			MyIntProp: 0,
			MyColorEnumProp: ColorEnum.NewEnumValue(),
		}
	}

	//...

	func NewManager() *boardgame.GameManager {
		//...

		//NewComponentChest will call Finish() on our Enums
		chest := boardgame.NewComponentChest(Enums)

		//...
	}

*/
package enum

import (
	"encoding/json"
	"errors"
	"math"
	"strconv"
)

//EnumSet is a set of enums where each Enum's values are unique. Normally you
//will create one in your package, add enums to it during initalization, and
//then use it for all managers you create.
type Set struct {
	finished bool
	enums    map[string]*Enum
	//A map of which int goes to which Enum
	values map[int]*Enum
}

//Enum is a named set of values within a set. Get a new one with
//enumSet.Add().
type Enum struct {
	name         string
	values       map[int]string
	defaultValue int
}

//variable is the underlying type we'll return for both Value and Constant.
type variable struct {
	enum *Enum
	val  int
}

//Const is an instantiation of an Enum that cannot be changed. You retrieve it
//from enum.NewConst(val).
type Const interface {
	Enum() *Enum
	Value() int
	String() string
	Copy() Const
}

//Var is an instantiation of a value that must be set to a value in
//the given enum. You retrieve one from enum.NewVar().
type Var interface {
	Const
	SetValue(int) error
	CopyVar() Var
}

//NewSet returns a new Set. Generally you'll call this once in a
//package and create the set during initalization.
func NewSet() *Set {
	return &Set{
		false,
		make(map[string]*Enum),
		make(map[int]*Enum),
	}
}

//MustCombineSets wraps CombineEnumSets, but instead of erroring will
//panic. Useful for package-level declarations outside of init().
func MustCombineSets(sets ...*Set) *Set {
	result, err := CombineSets(sets...)
	if err != nil {
		panic("Couldn't combine sets: " + err.Error())
	}
	return result
}

//CombineSets returns a new EnumSet that contains all of the EnumSets
//combined into one. The individual enums will literally be the same as the
//enums from the provided sets, so enum equality will work. Generally the sets
//have to know about each other, otherwise they are liable to overlap, which
//will error.
func CombineSets(sets ...*Set) (*Set, error) {
	result := NewSet()
	for i, set := range sets {
		for _, enumName := range set.EnumNames() {
			enum := set.Enum(enumName)
			if err := result.addEnum(enumName, enum); err != nil {
				return nil, errors.New("Couldn't add the " + strconv.Itoa(i) + " enumset because " + enumName + " had error: " + err.Error())
			}
		}
	}
	return result, nil
}

//Finish finalizes an EnumSet so that no more enums may be added. After this
//is called it is safe to use this in a multi-threaded environment. Repeated
//calls do nothing. ComponenChest automatically calls Finish() on the set you
//pass it.
func (e *Set) Finish() {
	e.finished = true
}

//EnumNames returns a list of all names in the Enum.
func (e *Set) EnumNames() []string {
	var result []string
	for key, _ := range e.enums {
		result = append(result, key)
	}
	return result
}

//Returns the Enum with the given name. In general you keep a reference to the
//enum yourself, but this is useful for programatically enumerating the enums.
func (e *Set) Enum(name string) *Enum {
	return e.enums[name]
}

//Membership returns the enum that the given val is a member of.
func (e *Set) Membership(val int) *Enum {
	return e.values[val]
}

//MustAdd is like Add, but instead of an error it will panic if the enum
//cannot be added. This is useful for defining your enums at the package level
//outside of an init().
func (e *Set) MustAdd(enumName string, values map[int]string) *Enum {
	result, err := e.Add(enumName, values)

	if err != nil {
		panic("Couldn't add to enumset: " + err.Error())
	}

	return result
}

/*
Add ads an enum with the given name and values to the enum manager. Will error
if that name has already been added, or any of the int values has been used
for any other enum item already. This means that enums must be unique within a
manager. Check out the package doc for the idiomatic way to initalize enums.
*/
func (e *Set) Add(enumName string, values map[int]string) (*Enum, error) {

	if len(values) == 0 {
		return nil, errors.New("No values provided")
	}

	enum := &Enum{
		enumName,
		make(map[int]string),
		math.MaxInt64,
	}

	seenValues := make(map[string]bool)

	for v, s := range values {

		if seenValues[s] {
			return nil, errors.New("String " + s + " was not unique within enum " + enumName)
		}

		seenValues[s] = true

		enum.values[v] = s

		if v < enum.defaultValue {
			enum.defaultValue = v
		}

	}
	if err := e.addEnum(enumName, enum); err != nil {
		return nil, err
	}
	return enum, nil
}

func (e *Set) addEnum(enumName string, enum *Enum) error {

	if e.finished {
		return errors.New("The set has been finished so no more enums can be added")
	}

	if _, ok := e.enums[enumName]; ok {
		return errors.New("That enum name has already been provided")
	}

	for v, _ := range enum.values {
		if _, ok := e.values[v]; ok {
			//Already registered
			return errors.New("Value " + strconv.Itoa(v) + " was registered twice")
		}

		e.values[v] = enum
	}

	e.enums[enumName] = enum

	return nil
}

//DefaultValue returns the default value for this enum (the lowest valid value
//in it).
func (e *Enum) DefaultValue() int {
	return e.defaultValue
}

//Valid returns whether the given value is a valid member of this enum.
func (e *Enum) Valid(val int) bool {
	_, ok := e.values[val]
	return ok
}

//String returns the string value associated with the given value.
func (e *Enum) String(val int) string {
	return e.values[val]
}

//Name returns the name of this enum; if set is the set this enum is part of,
//set.Enum(enum.Name()) == enum will be true.
func (e *Enum) Name() string {
	return e.name
}

//ValueFromString returns the enum value that corresponds to the given string,
//or -1 if no value has that string.
func (e *Enum) ValueFromString(in string) int {
	for v, str := range e.values {
		if str == in {
			return v
		}
	}
	return -1
}

//Copy returns a copy of the Value, that is equivalent, but will not be
//locked.
func (e *variable) Copy() Const {
	return &variable{
		e.enum,
		e.val,
	}
}

func (e *variable) CopyVar() Var {
	return &variable{
		e.enum,
		e.val,
	}
}

//NewEnumValue returns a new EnumValue associated with this enum, set to the
//Enum's DefaultValue to start.
func (e *Enum) NewVar() Var {
	return &variable{
		e,
		e.DefaultValue(),
	}
}

//MustNewConst is like NewConst, but if it would have errored it panics
//instead. It's convenient for initial set up where the whole app should fail
//to startup if it can't be configured anyway, and dealing with errors would
//be a lot of boilerplate.
func (e *Enum) MustNewConst(val int) Const {
	result, err := e.NewConst(val)
	if err != nil {
		panic("Couldn't create constant: " + err.Error())
	}
	return result
}

//NewConstant returns an enum.Constant that is permanently set to the provided
//val. If that value is not valid for this enum, it will error.
func (e *Enum) NewConst(val int) (Const, error) {
	variable := e.NewVar()
	if err := variable.SetValue(val); err != nil {
		return nil, err
	}
	return variable, nil
}

//The enum marshals as the string value of the enum so it's more readable.
func (e *variable) MarshalJSON() ([]byte, error) {
	return json.Marshal(e.String())
}

//UnmarshalJSON expects the blob to be the string value. Will error if that
//doesn't correspond to a valid value for this enum.
func (e *variable) UnmarshalJSON(blob []byte) error {
	var str string
	if err := json.Unmarshal(blob, &str); err != nil {
		return err
	}
	val := e.enum.ValueFromString(str)
	if val == -1 {
		return errors.New("That string value had no enum in the value")
	}
	return e.SetValue(val)
}

func (e *variable) Enum() *Enum {
	return e.enum
}

func (e *variable) Value() int {
	return e.val
}

func (e *variable) String() string {
	return e.enum.String(e.val)
}

//SetValue changes the value. Returns true if successful. Will fail if the
//value is locked or the val you want to set is not a valid number for the
//enum this value is associated with.
func (e *variable) SetValue(val int) error {
	if !e.enum.Valid(val) {
		return errors.New("That value is not valid for this enum")
	}
	e.val = val
	return nil
}
<|endoftext|>