<commit_before>// Copyright (c) 2020 Mercari, 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 internal

import (
	"fmt"
	"os"
	"sort"
	"strings"

	"github.com/kenshaw/snaker"
	"go.mercari.io/yo/models"
	"gopkg.in/yaml.v2"
)

type loaderImpl interface {
	ParamN(int) string
	MaskFunc() string
	ParseType(string, bool) (int, string, string)
	ValidCustomType(string, string) bool
	TableList() ([]*models.Table, error)
	ColumnList(string) ([]*models.Column, error)
	IndexList(string) ([]*models.Index, error)
	IndexColumnList(string, string) ([]*models.IndexColumn, error)
}

func NewTypeLoader(l loaderImpl, i Inflector) *TypeLoader {
	return &TypeLoader{loader: l, inflector: i}
}

// TypeLoader provides a common Loader implementation used by the built in
// schema/query loaders.
type TypeLoader struct {
	CustomTypes *models.CustomTypes
	loader      loaderImpl
	inflector   Inflector
}

// NthParam satisifies Loader's NthParam.
func (tl *TypeLoader) NthParam(i int) string {
	return tl.loader.ParamN(i)
}

// Mask returns the parameter mask.
func (tl *TypeLoader) Mask() string {

	return tl.loader.MaskFunc()
}

// LoadSchema loads schema definitions.
func (tl *TypeLoader) LoadSchema(args *ArgType) (map[string]*Type, map[string]*Index, error) {
	var err error

	// load tables
	tableMap, err := tl.LoadTable(args)
	if err != nil {
		return nil, nil, err
	}

	// load indexes
	ixMap, err := tl.LoadIndexes(args, tableMap)
	if err != nil {
		return nil, nil, err
	}

	setIndexesToTables(tableMap, ixMap)

	return tableMap, ixMap, nil
}

// LoadTable loads a schema table/view definition.
func (tl *TypeLoader) LoadTable(args *ArgType) (map[string]*Type, error) {
	var err error

	// load tables
	tableList, err := tl.loader.TableList()
	if err != nil {
		return nil, err
	}

	// tables
	tableMap := make(map[string]*Type)
	for _, ti := range tableList {
		ignore := false

		for _, ignoreTable := range args.IgnoreTables {
			if ignoreTable == ti.TableName {
				// Skip adding this table if user has specified they are not
				// interested.
				//
				// This could be useful for tables which are managed by the
				// database (e.g. SchemaMigrations) instead of
				// via Go code.
				ignore = true
			}
		}

		if ignore {
			continue
		}

		// create template
		typeTpl := &Type{
			Name:   SingularizeIdentifier(tl.inflector, ti.TableName),
			Schema: "",
			Fields: []*Field{},
			Table:  ti,
		}

		// process columns
		err = tl.LoadColumns(args, typeTpl)
		if err != nil {
			return nil, err
		}

		if err := tl.loadPrimaryKeys(typeTpl); err != nil {
			return nil, err
		}

		tableMap[ti.TableName] = typeTpl
	}

	return tableMap, nil
}

// loadPrimaryKeys loads primary key fields
func (tl *TypeLoader) loadPrimaryKeys(typeTpl *Type) error {
	// reorder primary keys
	indexCols, err := tl.loader.IndexColumnList(typeTpl.Table.TableName, "PRIMARY_KEY")
	if err != nil {
		panic(err)
	}

	var fields []*Field
	for _, idx := range indexCols {
		var field *Field
		for _, f := range typeTpl.Fields {
			if f.Col.ColumnName == idx.ColumnName {
				field = f
				break
			}
		}

		if field == nil {
			return fmt.Errorf("primary key column is not found in column list: table=%v column=%v",
				typeTpl.Name, idx.ColumnName,
			)
		}
		fields = append(fields, field)
	}

	typeTpl.PrimaryKey = fields[0] // backward compatibility
	typeTpl.PrimaryKeyFields = fields
	return nil
}

// tableCustomTypes find custom type definitions of the table
func (tl *TypeLoader) tableCustomTypes(table string) map[string]string {
	var columnTypes map[string]string
	if tl.CustomTypes != nil {
		for _, v := range tl.CustomTypes.Tables {
			if v.Name == table {
				columnTypes = v.Columns
				break
			}
		}
	}

	return columnTypes
}

// LoadColumns loads schema table/view columns.
func (tl *TypeLoader) LoadColumns(args *ArgType, typeTpl *Type) error {
	var err error

	// load columns
	columnList, err := tl.loader.ColumnList(typeTpl.Table.TableName)
	if err != nil {
		return err
	}

	columnTypes := tl.tableCustomTypes(typeTpl.Table.TableName)
	// process columns
	for _, c := range columnList {
		ignore := false

		for _, ignoreField := range args.IgnoreFields {
			if ignoreField == c.ColumnName {
				// Skip adding this field if user has specified they are not
				// interested.
				//
				// This could be useful for fields which are managed by the
				// database (e.g. automatically updated timestamps) instead of
				// via Go code.
				ignore = true
			}
		}

		if ignore {
			continue
		}

		// set col info
		f := &Field{
			Name: snaker.ForceCamelIdentifier(c.ColumnName),
			// Name: c.ColumnName,
			Col: c,
		}

		f.Len, f.NilType, f.Type = tl.loader.ParseType(c.DataType, !c.NotNull)

		// set custom type
		if columnTypes != nil {
			if t, ok := columnTypes[c.ColumnName]; ok && tl.loader.ValidCustomType(c.DataType, t) {
				f.CustomType = t
			}
		}

		// append col to template fields
		typeTpl.Fields = append(typeTpl.Fields, f)
	}

	return nil
}

// LoadIndexes loads schema index definitions.
func (tl *TypeLoader) LoadIndexes(args *ArgType, tableMap map[string]*Type) (map[string]*Index, error) {
	var err error

	ixMap := map[string]*Index{}
	for _, t := range tableMap {
		// load table indexes
		err = tl.LoadTableIndexes(args, t, ixMap)
		if err != nil {
			return nil, err
		}
	}

	return ixMap, nil
}

// LoadTableIndexes loads schema index definitions per table.
func (tl *TypeLoader) LoadTableIndexes(args *ArgType, typeTpl *Type, ixMap map[string]*Index) error {
	var err error
	var priIxLoaded bool

	// load indexes
	indexList, err := tl.loader.IndexList(typeTpl.Table.TableName)
	if err != nil {
		return err
	}

	// process indexes
	for _, ix := range indexList {
		// save whether or not the primary key index was processed
		priIxLoaded = priIxLoaded || ix.IsPrimary || (ix.Origin == "pk")

		// create index template
		ixTpl := &Index{
			Schema: "",
			Type:   typeTpl,
			Fields: []*Field{},
			Index:  ix,
		}

		// load index columns
		err = tl.LoadIndexColumns(args, ixTpl)
		if err != nil {
			return err
		}

		// build func name
		ixTpl.FuncName = tl.buildIndexFuncName(ixTpl)

		ixMap[typeTpl.Table.TableName+"_"+ix.IndexName] = ixTpl
	}

	return nil
}

func (tl *TypeLoader) buildIndexFuncName(ixTpl *Index) string {
	// build func name
	funcName := ixTpl.Type.Name
	if !ixTpl.Index.IsUnique {
		funcName = tl.inflector.Pluralize(ixTpl.Type.Name)
	}
	funcName = funcName + "By"

	// add param names
	paramNames := make([]string, 0, len(ixTpl.Fields))
	for _, f := range ixTpl.StoringFields {
		paramNames = append(paramNames, f.Name)
	}
	for _, f := range ixTpl.Fields {
		paramNames = append(paramNames, f.Name)
	}

	return funcName + strings.Join(paramNames, "")
}

// LoadIndexColumns loads the index column information.
func (tl *TypeLoader) LoadIndexColumns(args *ArgType, ixTpl *Index) error {
	var err error

	// load index columns
	indexCols, err := tl.loader.IndexColumnList(ixTpl.Type.Table.TableName, ixTpl.Index.IndexName)
	if err != nil {
		return err
	}

	// process index columns
	for _, ic := range indexCols {
		var field *Field

	fieldLoop:
		// find field
		for _, f := range ixTpl.Type.Fields {
			if f.Col.ColumnName == ic.ColumnName {
				field = f
				break fieldLoop
			}
		}

		if field == nil {
			continue
		}

		if ic.Storing {
			// Storing column is added to StoringFields
			ixTpl.StoringFields = append(ixTpl.StoringFields, field)
		} else {
			ixTpl.Fields = append(ixTpl.Fields, field)
		}
		if !field.Col.NotNull {
			ixTpl.NullableFields = append(ixTpl.NullableFields, field)
		}
	}

	return nil
}

// LoadCustomTypes loads custom type definition from yml file
func (tl *TypeLoader) LoadCustomTypes(path string) error {
	file, err := os.Open(path)
	if err != nil {
		return err
	}

	var ctypes models.CustomTypes
	if err := yaml.NewDecoder(file).Decode(&ctypes); err != nil {
		return err
	}

	tl.CustomTypes = &ctypes

	return nil
}

func setIndexesToTables(tableMap map[string]*Type, ixMap map[string]*Index) {
	indexes := make([]*Index, 0, len(ixMap))
	for _, ix := range ixMap {
		indexes = append(indexes, ix)
	}
	sort.Slice(indexes, func(i, j int) bool {
		return indexes[i].FuncName < indexes[j].FuncName
	})
	for tbl, t := range tableMap {
		for _, ix := range indexes {
			if ix.Type.Table.TableName == tbl {
				t.Indexes = append(t.Indexes, ix)
			}
		}
	}
}
<commit_msg>support ignore fields table.column format (#64)<commit_after>// Copyright (c) 2020 Mercari, 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 internal

import (
	"fmt"
	"os"
	"sort"
	"strings"

	"github.com/kenshaw/snaker"
	"go.mercari.io/yo/models"
	"gopkg.in/yaml.v2"
)

type loaderImpl interface {
	ParamN(int) string
	MaskFunc() string
	ParseType(string, bool) (int, string, string)
	ValidCustomType(string, string) bool
	TableList() ([]*models.Table, error)
	ColumnList(string) ([]*models.Column, error)
	IndexList(string) ([]*models.Index, error)
	IndexColumnList(string, string) ([]*models.IndexColumn, error)
}

func NewTypeLoader(l loaderImpl, i Inflector) *TypeLoader {
	return &TypeLoader{loader: l, inflector: i}
}

// TypeLoader provides a common Loader implementation used by the built in
// schema/query loaders.
type TypeLoader struct {
	CustomTypes *models.CustomTypes
	loader      loaderImpl
	inflector   Inflector
}

// NthParam satisifies Loader's NthParam.
func (tl *TypeLoader) NthParam(i int) string {
	return tl.loader.ParamN(i)
}

// Mask returns the parameter mask.
func (tl *TypeLoader) Mask() string {

	return tl.loader.MaskFunc()
}

// LoadSchema loads schema definitions.
func (tl *TypeLoader) LoadSchema(args *ArgType) (map[string]*Type, map[string]*Index, error) {
	var err error

	// load tables
	tableMap, err := tl.LoadTable(args)
	if err != nil {
		return nil, nil, err
	}

	// load indexes
	ixMap, err := tl.LoadIndexes(args, tableMap)
	if err != nil {
		return nil, nil, err
	}

	setIndexesToTables(tableMap, ixMap)

	return tableMap, ixMap, nil
}

// LoadTable loads a schema table/view definition.
func (tl *TypeLoader) LoadTable(args *ArgType) (map[string]*Type, error) {
	var err error

	// load tables
	tableList, err := tl.loader.TableList()
	if err != nil {
		return nil, err
	}

	// tables
	tableMap := make(map[string]*Type)
	for _, ti := range tableList {
		ignore := false

		for _, ignoreTable := range args.IgnoreTables {
			if ignoreTable == ti.TableName {
				// Skip adding this table if user has specified they are not
				// interested.
				//
				// This could be useful for tables which are managed by the
				// database (e.g. SchemaMigrations) instead of
				// via Go code.
				ignore = true
			}
		}

		if ignore {
			continue
		}

		// create template
		typeTpl := &Type{
			Name:   SingularizeIdentifier(tl.inflector, ti.TableName),
			Schema: "",
			Fields: []*Field{},
			Table:  ti,
		}

		// process columns
		err = tl.LoadColumns(args, typeTpl)
		if err != nil {
			return nil, err
		}

		if err := tl.loadPrimaryKeys(typeTpl); err != nil {
			return nil, err
		}

		tableMap[ti.TableName] = typeTpl
	}

	return tableMap, nil
}

// loadPrimaryKeys loads primary key fields
func (tl *TypeLoader) loadPrimaryKeys(typeTpl *Type) error {
	// reorder primary keys
	indexCols, err := tl.loader.IndexColumnList(typeTpl.Table.TableName, "PRIMARY_KEY")
	if err != nil {
		panic(err)
	}

	var fields []*Field
	for _, idx := range indexCols {
		var field *Field
		for _, f := range typeTpl.Fields {
			if f.Col.ColumnName == idx.ColumnName {
				field = f
				break
			}
		}

		if field == nil {
			return fmt.Errorf("primary key column is not found in column list: table=%v column=%v",
				typeTpl.Name, idx.ColumnName,
			)
		}
		fields = append(fields, field)
	}

	typeTpl.PrimaryKey = fields[0] // backward compatibility
	typeTpl.PrimaryKeyFields = fields
	return nil
}

// tableCustomTypes find custom type definitions of the table
func (tl *TypeLoader) tableCustomTypes(table string) map[string]string {
	var columnTypes map[string]string
	if tl.CustomTypes != nil {
		for _, v := range tl.CustomTypes.Tables {
			if v.Name == table {
				columnTypes = v.Columns
				break
			}
		}
	}

	return columnTypes
}

// LoadColumns loads schema table/view columns.
func (tl *TypeLoader) LoadColumns(args *ArgType, typeTpl *Type) error {
	var err error

	// load columns
	columnList, err := tl.loader.ColumnList(typeTpl.Table.TableName)
	if err != nil {
		return err
	}

	columnTypes := tl.tableCustomTypes(typeTpl.Table.TableName)
	// process columns
	for _, c := range columnList {
		ignore := false

		for _, ignoreField := range args.IgnoreFields {
			if strings.Contains(ignoreField, ".") {
				s := strings.Split(ignoreField, ".")
				if s[0] == typeTpl.Table.TableName && s[1] == c.ColumnName {
					ignore = true
				}
			} else if ignoreField == c.ColumnName {
				// Skip adding this field if user has specified they are not
				// interested.
				//
				// This could be useful for fields which are managed by the
				// database (e.g. automatically updated timestamps) instead of
				// via Go code.
				ignore = true
			}
		}

		if ignore {
			continue
		}

		// set col info
		f := &Field{
			Name: snaker.ForceCamelIdentifier(c.ColumnName),
			// Name: c.ColumnName,
			Col: c,
		}

		f.Len, f.NilType, f.Type = tl.loader.ParseType(c.DataType, !c.NotNull)

		// set custom type
		if columnTypes != nil {
			if t, ok := columnTypes[c.ColumnName]; ok && tl.loader.ValidCustomType(c.DataType, t) {
				f.CustomType = t
			}
		}

		// append col to template fields
		typeTpl.Fields = append(typeTpl.Fields, f)
	}

	return nil
}

// LoadIndexes loads schema index definitions.
func (tl *TypeLoader) LoadIndexes(args *ArgType, tableMap map[string]*Type) (map[string]*Index, error) {
	var err error

	ixMap := map[string]*Index{}
	for _, t := range tableMap {
		// load table indexes
		err = tl.LoadTableIndexes(args, t, ixMap)
		if err != nil {
			return nil, err
		}
	}

	return ixMap, nil
}

// LoadTableIndexes loads schema index definitions per table.
func (tl *TypeLoader) LoadTableIndexes(args *ArgType, typeTpl *Type, ixMap map[string]*Index) error {
	var err error
	var priIxLoaded bool

	// load indexes
	indexList, err := tl.loader.IndexList(typeTpl.Table.TableName)
	if err != nil {
		return err
	}

	// process indexes
	for _, ix := range indexList {
		// save whether or not the primary key index was processed
		priIxLoaded = priIxLoaded || ix.IsPrimary || (ix.Origin == "pk")

		// create index template
		ixTpl := &Index{
			Schema: "",
			Type:   typeTpl,
			Fields: []*Field{},
			Index:  ix,
		}

		// load index columns
		err = tl.LoadIndexColumns(args, ixTpl)
		if err != nil {
			return err
		}

		// build func name
		ixTpl.FuncName = tl.buildIndexFuncName(ixTpl)

		ixMap[typeTpl.Table.TableName+"_"+ix.IndexName] = ixTpl
	}

	return nil
}

func (tl *TypeLoader) buildIndexFuncName(ixTpl *Index) string {
	// build func name
	funcName := ixTpl.Type.Name
	if !ixTpl.Index.IsUnique {
		funcName = tl.inflector.Pluralize(ixTpl.Type.Name)
	}
	funcName = funcName + "By"

	// add param names
	paramNames := make([]string, 0, len(ixTpl.Fields))
	for _, f := range ixTpl.StoringFields {
		paramNames = append(paramNames, f.Name)
	}
	for _, f := range ixTpl.Fields {
		paramNames = append(paramNames, f.Name)
	}

	return funcName + strings.Join(paramNames, "")
}

// LoadIndexColumns loads the index column information.
func (tl *TypeLoader) LoadIndexColumns(args *ArgType, ixTpl *Index) error {
	var err error

	// load index columns
	indexCols, err := tl.loader.IndexColumnList(ixTpl.Type.Table.TableName, ixTpl.Index.IndexName)
	if err != nil {
		return err
	}

	// process index columns
	for _, ic := range indexCols {
		var field *Field

	fieldLoop:
		// find field
		for _, f := range ixTpl.Type.Fields {
			if f.Col.ColumnName == ic.ColumnName {
				field = f
				break fieldLoop
			}
		}

		if field == nil {
			continue
		}

		if ic.Storing {
			// Storing column is added to StoringFields
			ixTpl.StoringFields = append(ixTpl.StoringFields, field)
		} else {
			ixTpl.Fields = append(ixTpl.Fields, field)
		}
		if !field.Col.NotNull {
			ixTpl.NullableFields = append(ixTpl.NullableFields, field)
		}
	}

	return nil
}

// LoadCustomTypes loads custom type definition from yml file
func (tl *TypeLoader) LoadCustomTypes(path string) error {
	file, err := os.Open(path)
	if err != nil {
		return err
	}

	var ctypes models.CustomTypes
	if err := yaml.NewDecoder(file).Decode(&ctypes); err != nil {
		return err
	}

	tl.CustomTypes = &ctypes

	return nil
}

func setIndexesToTables(tableMap map[string]*Type, ixMap map[string]*Index) {
	indexes := make([]*Index, 0, len(ixMap))
	for _, ix := range ixMap {
		indexes = append(indexes, ix)
	}
	sort.Slice(indexes, func(i, j int) bool {
		return indexes[i].FuncName < indexes[j].FuncName
	})
	for tbl, t := range tableMap {
		for _, ix := range indexes {
			if ix.Type.Table.TableName == tbl {
				t.Indexes = append(t.Indexes, ix)
			}
		}
	}
}
<|endoftext|>


<commit_before>// Package gmetric provides a client for the ganglia gmetric API.
package gmetric

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"net"
	"os"
	"time"
)

var (
	zeroByte   = []byte{byte(0)}
	errNoAddrs = errors.New("gmetric: no addrs provided")
	errNotOpen = errors.New("gmetric: client not opened")
	errNoName  = errors.New("gmetric: metric has no name")
)

type slopeType string

// The slope types supported by Ganglia.
const (
	SlopeZero     = slopeType("zero")
	SlopePositive = slopeType("positive")
	SlopeNegative = slopeType("negative")
	SlopeBoth     = slopeType("both")
)

func (s slopeType) value() uint32 {
	switch s {
	case SlopeZero:
		return 0
	case SlopePositive:
		return 1
	case SlopeNegative:
		return 2
	case SlopeBoth:
		return 3
	}
	return 4
}

type valueType string

// The value types supported by Ganglia.
const (
	ValueString  = valueType("string")
	ValueUint8   = valueType("uint8")
	ValueInt8    = valueType("int8")
	ValueUint16  = valueType("uint16")
	ValueInt16   = valueType("int16")
	ValueUint32  = valueType("uint32")
	ValueInt32   = valueType("int32")
	ValueFloat32 = valueType("float")
	ValueFloat64 = valueType("double")
)

// Represents a collection of errors.
type MultiError []error

// Returns a concatenation of all the contained errors.
func (m MultiError) Error() string {
	var buf bytes.Buffer
	buf.WriteString("gmetric: multi-error:")
	for _, e := range m {
		buf.WriteRune('\n')
		buf.WriteString(e.Error())
	}
	return buf.String()
}

// A Client represents a set of connections to write metrics to. The Client is
// itself a Writer which writes the given bytes to all open connections.
type Client struct {
	io.Writer

	// The target addresses or in gmond.conf parlance the udp_send_channels.
	Addr []net.Addr

	// The actual hostname for the machine. If empty the default will be set
	// based on os.Hostname.
	Host string

	// Optional spoof name for the machine. Since the default is reverse DNS this
	// allows for overriding the hostname to make it useful.
	Spoof string

	// Also known as TMax, it defines the max time interval between which the
	// daemon will expect updates. This should map to how often you publish the
	// metric.
	TickInterval time.Duration

	// Also known as DMax, it defines the lifetime for the metric. That is, once
	// the last received metric is older than the defined value it will become
	// eligible for garbage collection.
	Lifetime time.Duration

	conn []net.Conn
}

// Defines a Metric.
type Metric struct {
	// The name is used as the file name, and also the title unless one is
	// explicitly provided.
	Name string

	// The title is for human consumption and is shown atop the graph.
	Title string

	// Descriptions serve as documentation.
	Description string

	// The groups ensure your metric is kept alongside sibling metrics.
	Groups []string

	// The units are shown in the graph to provide context to the numbers.
	Units string

	// The actual hostname for the machine.
	Host string

	// Optional spoof name for the machine. Since the default is reverse DNS this
	// allows for overriding the hostname to make it useful.
	Spoof string

	// Defines the value type. You must specify one of the predefined constants.
	ValueType valueType

	// Defines the slope type. You must specify one of the predefined constants.
	Slope slopeType

	// Also known as TMax, it defines the max time interval between which the
	// daemon will expect updates. This should map to how often you publish the
	// metric.
	TickInterval time.Duration

	// Also known as DMax, it defines the lifetime for the metric. That is, once
	// the last received metric is older than the defined value it will become
	// eligible for garbage collection.
	Lifetime time.Duration
}

// Writes a metadata packet for the Metric.
func (m *Metric) writeMeta(c *Client, w io.Writer) (err error) {
	pw := &panickyWriter{Writer: w}
	defer func() {
		if r := recover(); r != nil {
			if r == errPanickyWriter {
				err = pw.Error
			} else {
				panic(r)
			}
		}
	}()

	writeUint32(pw, 128) // identifies meta
	m.writeHead(c, pw)
	writeString(pw, string(m.ValueType))
	writeString(pw, m.Name)
	writeString(pw, m.Units)
	writeUint32(pw, m.Slope.value())

	tick := m.TickInterval.Seconds()
	if tick == 0 {
		writeUint32(pw, uint32(c.TickInterval.Seconds()))
	} else {
		writeUint32(pw, uint32(tick))
	}

	life := m.Lifetime.Seconds()
	if life == 0 {
		writeUint32(pw, uint32(c.Lifetime.Seconds()))
	} else {
		writeUint32(pw, uint32(life))
	}

	var extras [][2]string
	if m.Title != "" {
		extras = append(extras, [2]string{"TITLE", m.Title})
	}
	if m.Description != "" {
		extras = append(extras, [2]string{"DESC", m.Description})
	}

	spoof := m.Spoof
	if spoof == "" {
		spoof = c.Spoof
	}
	if spoof != "" {
		extras = append(extras, [2]string{"SPOOF_HOST", spoof})
	}

	for _, group := range m.Groups {
		extras = append(extras, [2]string{"GROUP", group})
	}
	writeExtras(pw, extras)
	return
}

// Writes a value packet for the given value. The value will be encoded based
// on the configured ValueType.
func (m *Metric) writeValue(c *Client, w io.Writer, val interface{}) (err error) {
	pw := &panickyWriter{Writer: w}
	defer func() {
		if r := recover(); r != nil {
			if r == errPanickyWriter {
				err = pw.Error
			} else {
				panic(r)
			}
		}
	}()

	writeUint32(pw, 133) // identifies a value
	m.writeHead(c, pw)
	writeString(pw, "%s")
	writeString(pw, fmt.Sprint(val))
	return
}

func (m *Metric) writeHead(c *Client, w io.Writer) {
	host := m.Host
	if host == "" {
		host = c.Host
	}

	spoof := m.Spoof
	if spoof == "" {
		spoof = c.Spoof
	}

	hasSpoof := spoof != ""
	if hasSpoof {
		writeString(w, spoof)
	} else {
		writeString(w, host)
	}
	writeString(w, m.Name)
	if hasSpoof {
		writeUint32(w, 1)
	} else {
		writeUint32(w, 0)
	}
}

// Write the Metric metadata.
func (c *Client) WriteMeta(m *Metric) error {
	if c.Writer == nil {
		return errNotOpen
	}
	if m.Name == "" {
		return errNoName
	}
	var buf bytes.Buffer
	if err := m.writeMeta(c, &buf); err != nil {
		return err
	}
	if _, err := c.Write(buf.Bytes()); err != nil {
		return err
	}
	return nil
}

// Write a value for the Metric.
func (c *Client) WriteValue(m *Metric, val interface{}) error {
	if c.Writer == nil {
		return errNotOpen
	}
	if m.Name == "" {
		return errNoName
	}
	var buf bytes.Buffer
	if err := m.writeValue(c, &buf, val); err != nil {
		return err
	}
	if _, err := c.Write(buf.Bytes()); err != nil {
		return err
	}
	return nil
}

// Open the connections. If an error is returned it will be a MultiError.
func (c *Client) Open() error {
	if len(c.Addr) == 0 {
		return errNoAddrs
	}

	if c.Host == "" {
		c.Host, _ = os.Hostname()
	}

	var errs MultiError
	var writers []io.Writer
	for _, addr := range c.Addr {
		s, err := net.Dial(addr.Network(), addr.String())
		if err != nil {
			errs = append(errs, err)
			continue
		}
		c.conn = append(c.conn, s)
		writers = append(writers, s)
	}
	c.Writer = io.MultiWriter(writers...)

	if len(errs) == 0 {
		return nil
	}
	return errs
}

// Close the connections. If an error is returned it will be a MultiError.
func (c *Client) Close() error {
	if len(c.Addr) == 0 {
		return errNoAddrs
	}

	var errs MultiError
	for _, conn := range c.conn {
		if err := conn.Close(); err != nil {
			errs = append(errs, err)
		}
	}

	if len(errs) == 0 {
		return nil
	}
	return errs
}

func writeUint32(w io.Writer, val uint32) {
	w.Write([]byte{
		byte(val >> 24 & 0xff),
		byte(val >> 16 & 0xff),
		byte(val >> 8 & 0xff),
		byte(val & 0xff),
	})
}

func writeString(w io.Writer, val string) {
	l := uint32(len(val))
	writeUint32(w, l)
	fmt.Fprint(w, val)
	offset := l % 4
	if offset != 0 {
		for j := offset; j < 4; j++ {
			w.Write(zeroByte)
		}
	}
}

func writeExtras(w io.Writer, extras [][2]string) {
	writeUint32(w, uint32(len(extras)))
	for _, p := range extras {
		writeString(w, p[0])
		writeString(w, p[1])
	}
}

var errPanickyWriter = errors.New("panicky-writer sentinel")

// Panicky Writer provides a io.Writer that panics whenever the underlying
// writer returns an error. This allows for localized panic/recover for less
// verbose error handling internally. DO NOT expose this without a
// corresponding recover().
type panickyWriter struct {
	io.Writer
	Count int
	Error error
}

func (p *panickyWriter) Write(b []byte) (int, error) {
	n, err := p.Writer.Write(b)
	p.Count += n
	if err != nil {
		p.Error = err
		panic(errPanickyWriter)
	}
	return n, nil
}
<commit_msg>writeCheck for shared checks<commit_after>// Package gmetric provides a client for the ganglia gmetric API.
package gmetric

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"net"
	"os"
	"time"
)

var (
	zeroByte   = []byte{byte(0)}
	errNoAddrs = errors.New("gmetric: no addrs provided")
	errNotOpen = errors.New("gmetric: client not opened")
	errNoName  = errors.New("gmetric: metric has no name")
)

type slopeType string

// The slope types supported by Ganglia.
const (
	SlopeZero     = slopeType("zero")
	SlopePositive = slopeType("positive")
	SlopeNegative = slopeType("negative")
	SlopeBoth     = slopeType("both")
)

func (s slopeType) value() uint32 {
	switch s {
	case SlopeZero:
		return 0
	case SlopePositive:
		return 1
	case SlopeNegative:
		return 2
	case SlopeBoth:
		return 3
	}
	return 4
}

type valueType string

// The value types supported by Ganglia.
const (
	ValueString  = valueType("string")
	ValueUint8   = valueType("uint8")
	ValueInt8    = valueType("int8")
	ValueUint16  = valueType("uint16")
	ValueInt16   = valueType("int16")
	ValueUint32  = valueType("uint32")
	ValueInt32   = valueType("int32")
	ValueFloat32 = valueType("float")
	ValueFloat64 = valueType("double")
)

// Represents a collection of errors.
type MultiError []error

// Returns a concatenation of all the contained errors.
func (m MultiError) Error() string {
	var buf bytes.Buffer
	buf.WriteString("gmetric: multi-error:")
	for _, e := range m {
		buf.WriteRune('\n')
		buf.WriteString(e.Error())
	}
	return buf.String()
}

// A Client represents a set of connections to write metrics to. The Client is
// itself a Writer which writes the given bytes to all open connections.
type Client struct {
	io.Writer

	// The target addresses or in gmond.conf parlance the udp_send_channels.
	Addr []net.Addr

	// The actual hostname for the machine. If empty the default will be set
	// based on os.Hostname.
	Host string

	// Optional spoof name for the machine. Since the default is reverse DNS this
	// allows for overriding the hostname to make it useful.
	Spoof string

	// Also known as TMax, it defines the max time interval between which the
	// daemon will expect updates. This should map to how often you publish the
	// metric.
	TickInterval time.Duration

	// Also known as DMax, it defines the lifetime for the metric. That is, once
	// the last received metric is older than the defined value it will become
	// eligible for garbage collection.
	Lifetime time.Duration

	conn []net.Conn
}

// Defines a Metric.
type Metric struct {
	// The name is used as the file name, and also the title unless one is
	// explicitly provided.
	Name string

	// The title is for human consumption and is shown atop the graph.
	Title string

	// Descriptions serve as documentation.
	Description string

	// The groups ensure your metric is kept alongside sibling metrics.
	Groups []string

	// The units are shown in the graph to provide context to the numbers.
	Units string

	// The actual hostname for the machine.
	Host string

	// Optional spoof name for the machine. Since the default is reverse DNS this
	// allows for overriding the hostname to make it useful.
	Spoof string

	// Defines the value type. You must specify one of the predefined constants.
	ValueType valueType

	// Defines the slope type. You must specify one of the predefined constants.
	Slope slopeType

	// Also known as TMax, it defines the max time interval between which the
	// daemon will expect updates. This should map to how often you publish the
	// metric.
	TickInterval time.Duration

	// Also known as DMax, it defines the lifetime for the metric. That is, once
	// the last received metric is older than the defined value it will become
	// eligible for garbage collection.
	Lifetime time.Duration
}

// Writes a metadata packet for the Metric.
func (m *Metric) writeMeta(c *Client, w io.Writer) (err error) {
	pw := &panickyWriter{Writer: w}
	defer func() {
		if r := recover(); r != nil {
			if r == errPanickyWriter {
				err = pw.Error
			} else {
				panic(r)
			}
		}
	}()

	writeUint32(pw, 128) // identifies meta
	m.writeHead(c, pw)
	writeString(pw, string(m.ValueType))
	writeString(pw, m.Name)
	writeString(pw, m.Units)
	writeUint32(pw, m.Slope.value())

	tick := m.TickInterval.Seconds()
	if tick == 0 {
		writeUint32(pw, uint32(c.TickInterval.Seconds()))
	} else {
		writeUint32(pw, uint32(tick))
	}

	life := m.Lifetime.Seconds()
	if life == 0 {
		writeUint32(pw, uint32(c.Lifetime.Seconds()))
	} else {
		writeUint32(pw, uint32(life))
	}

	var extras [][2]string
	if m.Title != "" {
		extras = append(extras, [2]string{"TITLE", m.Title})
	}
	if m.Description != "" {
		extras = append(extras, [2]string{"DESC", m.Description})
	}

	spoof := m.Spoof
	if spoof == "" {
		spoof = c.Spoof
	}
	if spoof != "" {
		extras = append(extras, [2]string{"SPOOF_HOST", spoof})
	}

	for _, group := range m.Groups {
		extras = append(extras, [2]string{"GROUP", group})
	}
	writeExtras(pw, extras)
	return
}

// Writes a value packet for the given value. The value will be encoded based
// on the configured ValueType.
func (m *Metric) writeValue(c *Client, w io.Writer, val interface{}) (err error) {
	pw := &panickyWriter{Writer: w}
	defer func() {
		if r := recover(); r != nil {
			if r == errPanickyWriter {
				err = pw.Error
			} else {
				panic(r)
			}
		}
	}()

	writeUint32(pw, 133) // identifies a value
	m.writeHead(c, pw)
	writeString(pw, "%s")
	writeString(pw, fmt.Sprint(val))
	return
}

func (m *Metric) writeHead(c *Client, w io.Writer) {
	host := m.Host
	if host == "" {
		host = c.Host
	}

	spoof := m.Spoof
	if spoof == "" {
		spoof = c.Spoof
	}

	hasSpoof := spoof != ""
	if hasSpoof {
		writeString(w, spoof)
	} else {
		writeString(w, host)
	}
	writeString(w, m.Name)
	if hasSpoof {
		writeUint32(w, 1)
	} else {
		writeUint32(w, 0)
	}
}

func (c *Client) writeCheck(m *Metric) error {
	if c.Writer == nil {
		return errNotOpen
	}
	if m.Name == "" {
		return errNoName
	}
	return nil
}

// Write the Metric metadata.
func (c *Client) WriteMeta(m *Metric) error {
	if err := c.writeCheck(m); err != nil {
		return err
	}
	var buf bytes.Buffer
	if err := m.writeMeta(c, &buf); err != nil {
		return err
	}
	if _, err := c.Write(buf.Bytes()); err != nil {
		return err
	}
	return nil
}

// Write a value for the Metric.
func (c *Client) WriteValue(m *Metric, val interface{}) error {
	if err := c.writeCheck(m); err != nil {
		return err
	}
	var buf bytes.Buffer
	if err := m.writeValue(c, &buf, val); err != nil {
		return err
	}
	if _, err := c.Write(buf.Bytes()); err != nil {
		return err
	}
	return nil
}

// Open the connections. If an error is returned it will be a MultiError.
func (c *Client) Open() error {
	if len(c.Addr) == 0 {
		return errNoAddrs
	}

	if c.Host == "" {
		c.Host, _ = os.Hostname()
	}

	var errs MultiError
	var writers []io.Writer
	for _, addr := range c.Addr {
		s, err := net.Dial(addr.Network(), addr.String())
		if err != nil {
			errs = append(errs, err)
			continue
		}
		c.conn = append(c.conn, s)
		writers = append(writers, s)
	}
	c.Writer = io.MultiWriter(writers...)

	if len(errs) == 0 {
		return nil
	}
	return errs
}

// Close the connections. If an error is returned it will be a MultiError.
func (c *Client) Close() error {
	if len(c.Addr) == 0 {
		return errNoAddrs
	}

	var errs MultiError
	for _, conn := range c.conn {
		if err := conn.Close(); err != nil {
			errs = append(errs, err)
		}
	}

	if len(errs) == 0 {
		return nil
	}
	return errs
}

func writeUint32(w io.Writer, val uint32) {
	w.Write([]byte{
		byte(val >> 24 & 0xff),
		byte(val >> 16 & 0xff),
		byte(val >> 8 & 0xff),
		byte(val & 0xff),
	})
}

func writeString(w io.Writer, val string) {
	l := uint32(len(val))
	writeUint32(w, l)
	fmt.Fprint(w, val)
	offset := l % 4
	if offset != 0 {
		for j := offset; j < 4; j++ {
			w.Write(zeroByte)
		}
	}
}

func writeExtras(w io.Writer, extras [][2]string) {
	writeUint32(w, uint32(len(extras)))
	for _, p := range extras {
		writeString(w, p[0])
		writeString(w, p[1])
	}
}

var errPanickyWriter = errors.New("panicky-writer sentinel")

// Panicky Writer provides a io.Writer that panics whenever the underlying
// writer returns an error. This allows for localized panic/recover for less
// verbose error handling internally. DO NOT expose this without a
// corresponding recover().
type panickyWriter struct {
	io.Writer
	Count int
	Error error
}

func (p *panickyWriter) Write(b []byte) (int, error) {
	n, err := p.Writer.Write(b)
	p.Count += n
	if err != nil {
		p.Error = err
		panic(errPanickyWriter)
	}
	return n, nil
}
<|endoftext|>


<commit_before>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

// This is the main login engine.

package engine

import (
	"errors"
	"strings"

	"github.com/keybase/client/go/libkb"
	keybase1 "github.com/keybase/client/go/protocol/keybase1"
)

var errNoConfig = errors.New("No user config available")
var errNoDevice = errors.New("No device provisioned locally for this user")

// Login is an engine.
type Login struct {
	libkb.Contextified
	deviceType      string
	usernameOrEmail string
	clientType      keybase1.ClientType
}

// NewLogin creates a Login engine.  username is optional.
// deviceType should be libkb.DeviceTypeDesktop or
// libkb.DeviceTypeMobile.
func NewLogin(g *libkb.GlobalContext, deviceType string, usernameOrEmail string, ct keybase1.ClientType) *Login {
	return &Login{
		Contextified:    libkb.NewContextified(g),
		deviceType:      deviceType,
		usernameOrEmail: strings.TrimSpace(usernameOrEmail),
		clientType:      ct,
	}
}

// Name is the unique engine name.
func (e *Login) Name() string {
	return "Login"
}

// GetPrereqs returns the engine prereqs.
func (e *Login) Prereqs() Prereqs {
	return Prereqs{}
}

// RequiredUIs returns the required UIs.
func (e *Login) RequiredUIs() []libkb.UIKind {
	return []libkb.UIKind{}
}

// SubConsumers returns the other UI consumers for this engine.
func (e *Login) SubConsumers() []libkb.UIConsumer {
	return []libkb.UIConsumer{
		&LoginProvisionedDevice{},
		&loginLoadUser{},
		&loginProvision{},
	}
}

// Run starts the engine.
func (e *Login) Run(ctx *Context) error {
	// check to see if already logged in
	loggedInOK, err := e.checkLoggedIn(ctx)
	if err != nil {
		return err
	}
	if loggedInOK {
		return nil
	}
	e.G().Log.Debug("Login: not currently logged in")

	if len(e.usernameOrEmail) > 0 && libkb.CheckEmail.F(e.usernameOrEmail) {
		// If e.usernameOrEmail is provided and it is an email address, then
		// loginProvisionedDevice is pointless.  It would return an error,
		// but might as well not even use it.
		e.G().Log.Debug("skipping loginProvisionedDevice since %q provided to Login, which looks like an email address.", e.usernameOrEmail)
	} else {
		// First see if this device is already provisioned and it is possible to log in.
		loggedInOK, err := e.loginProvisionedDevice(ctx, e.usernameOrEmail)
		if err != nil {
			e.G().Log.Debug("loginProvisionedDevice error: %s", err)
			return err
		}
		if loggedInOK {
			e.G().Log.Debug("loginProvisionedDevice success")
			return nil
		}

		e.G().Log.Debug("loginProvisionedDevice failed, continuing with device provisioning")
	}

	// clear out any existing session:
	e.G().Log.Debug("clearing any exising login session with Logout before loading user for login")
	e.G().Logout()

	// run the LoginLoadUser sub-engine to load a user
	e.G().Log.Debug("loading login user for %q", e.usernameOrEmail)
	ueng := newLoginLoadUser(e.G(), e.usernameOrEmail)
	if err := RunEngine(ueng, ctx); err != nil {
		return err
	}

	// make sure the user isn't already provisioned (can
	// get here if usernameOrEmail is an email address
	// for an already provisioned on this device user).
	if ueng.User().HasCurrentDeviceInCurrentInstall() {
		e.G().Log.Debug("user %q (%s) has previously provisioned this device, trying to login on it", e.usernameOrEmail, ueng.User().GetName())
		loggedInOK, err := e.loginProvisionedDevice(ctx, ueng.User().GetName())
		if err != nil {
			e.G().Log.Debug("loginProvisionedDevice after loginLoadUser error: %s", err)
			return err
		}
		if loggedInOK {
			e.G().Log.Debug("loginProvisionedDevice after loginLoadUser success")
			return nil
		}

		// this shouldn't happen:
		e.G().Log.Debug("loginProvisionedDevice after loginLoadUser (and user had current deivce in current install), failed to login [unexpected]")
		return libkb.DeviceAlreadyProvisionedError{}
	}

	e.G().Log.Debug("attempting device provisioning")

	darg := &loginProvisionArg{
		DeviceType: e.deviceType,
		ClientType: e.clientType,
		User:       ueng.User(),
	}
	deng := newLoginProvision(e.G(), darg)
	if err := RunEngine(deng, ctx); err != nil {
		return err
	}

	e.perUserKeyUpgradeSoft(ctx)

	e.G().Log.Debug("Login provisioning success, sending login notification")
	e.sendNotification()
	return nil
}

// notProvisionedErr will return true if err signifies that login
// failed because this device has not yet been provisioned.
func (e *Login) notProvisionedErr(err error) bool {
	if err == errNoDevice {
		return true
	}
	if err == errNoConfig {
		return true
	}

	e.G().Log.Debug("notProvisioned, not handling error %s (err type: %T)", err, err)

	return false
}

func (e *Login) sendNotification() {
	e.G().NotifyRouter.HandleLogin(string(e.G().Env.GetUsername()))
	e.G().CallLoginHooks()
}

// Get a per-user key.
// Wait for attempt but only warn on error.
func (e *Login) perUserKeyUpgradeSoft(ctx *Context) error {
	eng := NewPerUserKeyUpgrade(e.G(), &PerUserKeyUpgradeArgs{})
	err := RunEngine(eng, ctx)
	if err != nil {
		e.G().Log.CWarningf(ctx.GetNetContext(), "loginProvision PerUserKeyUpgrade failed: %v", err)
	}
	return err
}

func (e *Login) checkLoggedIn(ctx *Context) (bool, error) {
	if !e.G().ActiveDevice.Valid() {
		return false, nil
	}

	if len(e.usernameOrEmail) == 0 {
		e.G().Log.Debug("Login: already logged in, no username or email provided, so returning without error")
		return true, nil
	}
	if libkb.CheckEmail.F(e.usernameOrEmail) {
		e.G().Log.Debug("Login: already logged in, but %q email address provided.  Can't determine if that is current user without further work, so just returning LoggedInError")
		return true, libkb.LoggedInError{}
	}
	username, err := e.G().GetUPAKLoader().LookupUsername(ctx.NetContext, e.G().ActiveDevice.UID())
	if err != nil {
		return true, err
	}
	if username.Eq(libkb.NewNormalizedUsername(e.usernameOrEmail)) {
		e.G().Log.Debug("Login: already logged in as %q, returning without error", e.usernameOrEmail)
		return true, nil
	}

	e.G().Log.Debug("Login: logged in already as %q (%q requested), returning LoggedInError", username, e.usernameOrEmail)
	return true, libkb.LoggedInError{}

}

func (e *Login) loginProvisionedDevice(ctx *Context, username string) (bool, error) {
	eng := NewLoginProvisionedDevice(e.G(), username)
	err := RunEngine(eng, ctx)
	if err == nil {
		// login successful
		e.G().Log.Debug("LoginProvisionedDevice.Run() was successful")
		// Note:  LoginProvisionedDevice Run() will send login notifications, no need to
		// send here.
		return true, nil
	}

	// if this device has been provisioned already and there was an error, then
	// return that error.  Otherwise, ignore it and keep going.
	if !e.notProvisionedErr(err) {
		return false, err
	}

	e.G().Log.Debug("loginProvisionedDevice error: %s (continuing with device provisioning...)", err)

	return false, nil
}
<commit_msg>Clarify log message<commit_after>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

// This is the main login engine.

package engine

import (
	"errors"
	"strings"

	"github.com/keybase/client/go/libkb"
	keybase1 "github.com/keybase/client/go/protocol/keybase1"
)

var errNoConfig = errors.New("No user config available")
var errNoDevice = errors.New("No device provisioned locally for this user")

// Login is an engine.
type Login struct {
	libkb.Contextified
	deviceType      string
	usernameOrEmail string
	clientType      keybase1.ClientType
}

// NewLogin creates a Login engine.  username is optional.
// deviceType should be libkb.DeviceTypeDesktop or
// libkb.DeviceTypeMobile.
func NewLogin(g *libkb.GlobalContext, deviceType string, usernameOrEmail string, ct keybase1.ClientType) *Login {
	return &Login{
		Contextified:    libkb.NewContextified(g),
		deviceType:      deviceType,
		usernameOrEmail: strings.TrimSpace(usernameOrEmail),
		clientType:      ct,
	}
}

// Name is the unique engine name.
func (e *Login) Name() string {
	return "Login"
}

// GetPrereqs returns the engine prereqs.
func (e *Login) Prereqs() Prereqs {
	return Prereqs{}
}

// RequiredUIs returns the required UIs.
func (e *Login) RequiredUIs() []libkb.UIKind {
	return []libkb.UIKind{}
}

// SubConsumers returns the other UI consumers for this engine.
func (e *Login) SubConsumers() []libkb.UIConsumer {
	return []libkb.UIConsumer{
		&LoginProvisionedDevice{},
		&loginLoadUser{},
		&loginProvision{},
	}
}

// Run starts the engine.
func (e *Login) Run(ctx *Context) error {
	// check to see if already logged in
	loggedInOK, err := e.checkLoggedIn(ctx)
	if err != nil {
		return err
	}
	if loggedInOK {
		return nil
	}
	e.G().Log.Debug("Login: not currently logged in")

	if len(e.usernameOrEmail) > 0 && libkb.CheckEmail.F(e.usernameOrEmail) {
		// If e.usernameOrEmail is provided and it is an email address, then
		// loginProvisionedDevice is pointless.  It would return an error,
		// but might as well not even use it.
		e.G().Log.Debug("skipping loginProvisionedDevice since %q provided to Login, which looks like an email address.", e.usernameOrEmail)
	} else {
		// First see if this device is already provisioned and it is possible to log in.
		loggedInOK, err := e.loginProvisionedDevice(ctx, e.usernameOrEmail)
		if err != nil {
			e.G().Log.Debug("loginProvisionedDevice error: %s", err)
			return err
		}
		if loggedInOK {
			e.G().Log.Debug("loginProvisionedDevice success")
			return nil
		}

		e.G().Log.Debug("loginProvisionedDevice failed, continuing with device provisioning")
	}

	// clear out any existing session:
	e.G().Log.Debug("clearing any exising login session with Logout before loading user for login")
	e.G().Logout()

	// run the LoginLoadUser sub-engine to load a user
	e.G().Log.Debug("loading login user for %q", e.usernameOrEmail)
	ueng := newLoginLoadUser(e.G(), e.usernameOrEmail)
	if err := RunEngine(ueng, ctx); err != nil {
		return err
	}

	// make sure the user isn't already provisioned (can
	// get here if usernameOrEmail is an email address
	// for an already provisioned on this device user).
	if ueng.User().HasCurrentDeviceInCurrentInstall() {
		e.G().Log.Debug("user %q (%s) has previously provisioned this device, trying to login on it", e.usernameOrEmail, ueng.User().GetName())
		loggedInOK, err := e.loginProvisionedDevice(ctx, ueng.User().GetName())
		if err != nil {
			e.G().Log.Debug("loginProvisionedDevice after loginLoadUser error: %s", err)
			return err
		}
		if loggedInOK {
			e.G().Log.Debug("loginProvisionedDevice after loginLoadUser success")
			return nil
		}

		// this shouldn't happen:
		e.G().Log.Debug("loginProvisionedDevice after loginLoadUser (and user had current deivce in current install), failed to login [unexpected]")
		return libkb.DeviceAlreadyProvisionedError{}
	}

	e.G().Log.Debug("attempting device provisioning")

	darg := &loginProvisionArg{
		DeviceType: e.deviceType,
		ClientType: e.clientType,
		User:       ueng.User(),
	}
	deng := newLoginProvision(e.G(), darg)
	if err := RunEngine(deng, ctx); err != nil {
		return err
	}

	e.perUserKeyUpgradeSoft(ctx)

	e.G().Log.Debug("Login provisioning success, sending login notification")
	e.sendNotification()
	return nil
}

// notProvisionedErr will return true if err signifies that login
// failed because this device has not yet been provisioned.
func (e *Login) notProvisionedErr(err error) bool {
	if err == errNoDevice {
		return true
	}
	if err == errNoConfig {
		return true
	}

	e.G().Log.Debug("notProvisioned, not handling error %s (err type: %T)", err, err)

	return false
}

func (e *Login) sendNotification() {
	e.G().NotifyRouter.HandleLogin(string(e.G().Env.GetUsername()))
	e.G().CallLoginHooks()
}

// Get a per-user key.
// Wait for attempt but only warn on error.
func (e *Login) perUserKeyUpgradeSoft(ctx *Context) error {
	eng := NewPerUserKeyUpgrade(e.G(), &PerUserKeyUpgradeArgs{})
	err := RunEngine(eng, ctx)
	if err != nil {
		e.G().Log.CWarningf(ctx.GetNetContext(), "loginProvision PerUserKeyUpgrade failed: %v", err)
	}
	return err
}

func (e *Login) checkLoggedIn(ctx *Context) (bool, error) {
	if !e.G().ActiveDevice.Valid() {
		return false, nil
	}

	if len(e.usernameOrEmail) == 0 {
		e.G().Log.Debug("Login: already logged in, no username or email provided, so returning without error")
		return true, nil
	}
	if libkb.CheckEmail.F(e.usernameOrEmail) {
		e.G().Log.Debug("Login: already logged in, but %q email address provided.  Can't determine if that is current user without further work, so just returning LoggedInError")
		return true, libkb.LoggedInError{}
	}
	username, err := e.G().GetUPAKLoader().LookupUsername(ctx.NetContext, e.G().ActiveDevice.UID())
	if err != nil {
		return true, err
	}
	if username.Eq(libkb.NewNormalizedUsername(e.usernameOrEmail)) {
		e.G().Log.Debug("Login: already logged in as %q, returning without error", e.usernameOrEmail)
		return true, nil
	}

	e.G().Log.Debug("Login: logged in already as %q (%q requested), returning LoggedInError", username, e.usernameOrEmail)
	return true, libkb.LoggedInError{}

}

func (e *Login) loginProvisionedDevice(ctx *Context, username string) (bool, error) {
	eng := NewLoginProvisionedDevice(e.G(), username)
	err := RunEngine(eng, ctx)
	if err == nil {
		// login successful
		e.G().Log.Debug("LoginProvisionedDevice.Run() was successful")
		// Note:  LoginProvisionedDevice Run() will send login notifications, no need to
		// send here.
		return true, nil
	}

	// if this device has been provisioned already and there was an error, then
	// return that error.  Otherwise, ignore it and keep going.
	if !e.notProvisionedErr(err) {
		return false, err
	}

	e.G().Log.Debug("loginProvisionedDevice error: %s (can continue to provision this device)", err)

	return false, nil
}
<|endoftext|>


<commit_before>package models

import (
	"bytes"
	"encoding/hex"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"regexp"
	"strings"
	"sync"

	"github.com/bnagy/gapstone"
	ks "github.com/keystone-engine/keystone/bindings/go/keystone"
)

var demangleRe = regexp.MustCompile(`^[^(]+`)

func Demangle(name string) string {
	if strings.HasPrefix(name, "__Z") {
		name = name[1:]
	} else if !strings.HasPrefix(name, "_Z") {
		return name
	}
	cmd := exec.Command("c++filt", "-n")
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return name
	}
	stdin, err := cmd.StdinPipe()
	if err != nil {
		return name
	}
	if err = cmd.Start(); err != nil {
		fmt.Fprintf(os.Stderr, "cmd error", err)
		return name
	}
	stdin.Write([]byte(name + "\n"))
	stdin.Close()
	out, err := ioutil.ReadAll(stdout)
	out = bytes.Trim(out, "\t\r\n ")
	if err != nil || len(out) == 0 {
		return name
	}
	cmd.Wait()
	out = demangleRe.FindSubmatch(out)[0]
	return string(out)
}

func Assemble(asm string, addr uint64, arch *Arch) ([]byte, error) {
	if arch.ks == nil {
		engine, err := ks.New(arch.KS_ARCH, arch.KS_MODE)
		if err != nil {
			return nil, err
		}
		arch.ks = engine
	}
	out, _, ok := arch.ks.Assemble(asm, addr)
	if !ok {
		return nil, arch.ks.LastError()
	}
	return out, nil
}

var discache = make(map[string]string)
var discacheLock sync.RWMutex

func Disas(mem []byte, addr uint64, arch *Arch, pad ...int) (string, error) {
	var asm []gapstone.Instruction
	cacheKey := fmt.Sprintf("%d|%s", addr, mem)
	if len(mem) == 0 {
		return "", nil
	}
	discacheLock.RLock()
	if cached, ok := discache[cacheKey]; ok {
		discacheLock.RUnlock()
		return cached, nil
	}
	discacheLock.RUnlock()
	if arch.cs == nil {
		engine, err := gapstone.New(arch.CS_ARCH, arch.CS_MODE)
		if err != nil {
			return "", err
		}
		arch.cs = &engine
	}
	asm, err := arch.cs.Disasm(mem, addr, 0)
	if err != nil {
		return "", err
	}
	var width uint
	if len(pad) > 0 {
		width = uint(pad[0])
	}
	for _, insn := range asm {
		if insn.Size > width {
			width = insn.Size
		}
	}
	var out []string
	for _, insn := range asm {
		pad := strings.Repeat(" ", int(width-insn.Size)*2)
		data := pad + hex.EncodeToString(insn.Bytes)
		out = append(out, fmt.Sprintf("0x%x: %s %s %s", insn.Address, data, insn.Mnemonic, insn.OpStr))
	}
	ret := strings.Join(out, "\n")
	discacheLock.Lock()
	discache[cacheKey] = ret
	discacheLock.Unlock()
	return ret, nil
}

func Repr(p []byte, strsize int) string {
	tmp := make([]string, len(p))
	for i, b := range p {
		if b >= 0x20 && b <= 0x7e {
			tmp[i] = string(b)
		} else {
			tmp[i] = fmt.Sprintf("\\x%02x", b)
		}
	}
	out := strings.Join(tmp, "")
	if strsize > 0 && len(out) > strsize {
		for i := len(tmp) - 1; len(out) > strsize-3; i-- {
			out = strings.Join(tmp[:i], "")
		}
		return "\"" + out + "\"..."
	}
	return "\"" + out + "\""
}

func HexDump(base uint64, mem []byte, bits int) []string {
	var clean = func(p []byte) string {
		o := make([]byte, len(p))
		for i, c := range p {
			if c >= 0x20 && c <= 0x7e {
				o[i] = c
			} else {
				o[i] = '.'
			}
		}
		return string(o)
	}
	bsz := bits / 8
	hexFmt := fmt.Sprintf("0x%%0%dx:", bsz*2)
	padBlock := strings.Repeat(" ", bsz*2)
	padTail := strings.Repeat(" ", bsz)

	width := 80
	addrSize := bsz*2 + 4
	blockCount := ((width - addrSize) * 3 / 4) / ((bsz + 1) * 2)
	lineSize := blockCount * bsz
	var out []string
	blocks := make([]string, blockCount)
	tail := make([]string, blockCount)
	for i := 0; i < len(mem); i += lineSize {
		memLine := mem[i:]
		for j := 0; j < blockCount; j++ {
			if j*bsz < len(memLine) {
				end := (j + 1) * bsz
				var block []byte
				if end > len(memLine) {
					block = memLine[j*bsz:]
				} else {
					block = memLine[j*bsz : end]
				}
				blocks[j] = hex.EncodeToString(block)
				tail[j] = clean(block)
				// if block was too short, pad with spaces
				if end > len(memLine) {
					pad := end - len(memLine)
					blocks[j] += strings.Repeat("  ", pad)
					tail[j] += strings.Repeat(" ", pad)
				}
			} else {
				blocks[j] = padBlock
				tail[j] = padTail
			}
		}
		line := []string{fmt.Sprintf(hexFmt, base+uint64(i))}
		line = append(line, strings.Join(blocks, " "))
		line = append(line, fmt.Sprintf("[%s]", strings.Join(tail, " ")))
		out = append(out, strings.Join(line, " "))
	}
	return out
}
<commit_msg>improve repr<commit_after>package models

import (
	"bytes"
	"encoding/hex"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"regexp"
	"strings"
	"sync"

	"github.com/bnagy/gapstone"
	ks "github.com/keystone-engine/keystone/bindings/go/keystone"
)

var demangleRe = regexp.MustCompile(`^[^(]+`)

func Demangle(name string) string {
	if strings.HasPrefix(name, "__Z") {
		name = name[1:]
	} else if !strings.HasPrefix(name, "_Z") {
		return name
	}
	cmd := exec.Command("c++filt", "-n")
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return name
	}
	stdin, err := cmd.StdinPipe()
	if err != nil {
		return name
	}
	if err = cmd.Start(); err != nil {
		fmt.Fprintf(os.Stderr, "cmd error", err)
		return name
	}
	stdin.Write([]byte(name + "\n"))
	stdin.Close()
	out, err := ioutil.ReadAll(stdout)
	out = bytes.Trim(out, "\t\r\n ")
	if err != nil || len(out) == 0 {
		return name
	}
	cmd.Wait()
	out = demangleRe.FindSubmatch(out)[0]
	return string(out)
}

func Assemble(asm string, addr uint64, arch *Arch) ([]byte, error) {
	if arch.ks == nil {
		engine, err := ks.New(arch.KS_ARCH, arch.KS_MODE)
		if err != nil {
			return nil, err
		}
		arch.ks = engine
	}
	out, _, ok := arch.ks.Assemble(asm, addr)
	if !ok {
		return nil, arch.ks.LastError()
	}
	return out, nil
}

var discache = make(map[string]string)
var discacheLock sync.RWMutex

func Disas(mem []byte, addr uint64, arch *Arch, pad ...int) (string, error) {
	var asm []gapstone.Instruction
	cacheKey := fmt.Sprintf("%d|%s", addr, mem)
	if len(mem) == 0 {
		return "", nil
	}
	discacheLock.RLock()
	if cached, ok := discache[cacheKey]; ok {
		discacheLock.RUnlock()
		return cached, nil
	}
	discacheLock.RUnlock()
	if arch.cs == nil {
		engine, err := gapstone.New(arch.CS_ARCH, arch.CS_MODE)
		if err != nil {
			return "", err
		}
		arch.cs = &engine
	}
	asm, err := arch.cs.Disasm(mem, addr, 0)
	if err != nil {
		return "", err
	}
	var width uint
	if len(pad) > 0 {
		width = uint(pad[0])
	}
	for _, insn := range asm {
		if insn.Size > width {
			width = insn.Size
		}
	}
	var out []string
	for _, insn := range asm {
		pad := strings.Repeat(" ", int(width-insn.Size)*2)
		data := pad + hex.EncodeToString(insn.Bytes)
		out = append(out, fmt.Sprintf("0x%x: %s %s %s", insn.Address, data, insn.Mnemonic, insn.OpStr))
	}
	ret := strings.Join(out, "\n")
	discacheLock.Lock()
	discache[cacheKey] = ret
	discacheLock.Unlock()
	return ret, nil
}

func Repr(p []byte, strsize int) string {
	tmp := make([]string, len(p))
	for i, b := range p {
		if b >= 0x20 && b <= 0x7e {
			tmp[i] = string(b)
		} else {
			switch b {
			case 0:
				tmp[i] = "\\0"
			case '\b':
				tmp[i] = "\\b"
			case '\r':
				tmp[i] = "\\r"
			case '\n':
				tmp[i] = "\\n"
			case '\t':
				tmp[i] = "\\t"
			default:
				tmp[i] = fmt.Sprintf("\\x%02x", b)
			}
		}
	}
	out := strings.Join(tmp, "")
	if strsize > 0 && len(out) > strsize {
		for i := len(tmp) - 1; len(out) > strsize-3; i-- {
			out = strings.Join(tmp[:i], "")
		}
		return "\"" + out + "\"..."
	}
	return "\"" + out + "\""
}

func HexDump(base uint64, mem []byte, bits int) []string {
	var clean = func(p []byte) string {
		o := make([]byte, len(p))
		for i, c := range p {
			if c >= 0x20 && c <= 0x7e {
				o[i] = c
			} else {
				o[i] = '.'
			}
		}
		return string(o)
	}
	bsz := bits / 8
	hexFmt := fmt.Sprintf("0x%%0%dx:", bsz*2)
	padBlock := strings.Repeat(" ", bsz*2)
	padTail := strings.Repeat(" ", bsz)

	width := 80
	addrSize := bsz*2 + 4
	blockCount := ((width - addrSize) * 3 / 4) / ((bsz + 1) * 2)
	lineSize := blockCount * bsz
	var out []string
	blocks := make([]string, blockCount)
	tail := make([]string, blockCount)
	for i := 0; i < len(mem); i += lineSize {
		memLine := mem[i:]
		for j := 0; j < blockCount; j++ {
			if j*bsz < len(memLine) {
				end := (j + 1) * bsz
				var block []byte
				if end > len(memLine) {
					block = memLine[j*bsz:]
				} else {
					block = memLine[j*bsz : end]
				}
				blocks[j] = hex.EncodeToString(block)
				tail[j] = clean(block)
				// if block was too short, pad with spaces
				if end > len(memLine) {
					pad := end - len(memLine)
					blocks[j] += strings.Repeat("  ", pad)
					tail[j] += strings.Repeat(" ", pad)
				}
			} else {
				blocks[j] = padBlock
				tail[j] = padTail
			}
		}
		line := []string{fmt.Sprintf(hexFmt, base+uint64(i))}
		line = append(line, strings.Join(blocks, " "))
		line = append(line, fmt.Sprintf("[%s]", strings.Join(tail, " ")))
		out = append(out, strings.Join(line, " "))
	}
	return out
}
<|endoftext|>


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

package ssa

// CreateTestMainPackage synthesizes a main package that runs all the
// tests of the supplied packages.
// It is closely coupled to $GOROOT/src/cmd/go/test.go and $GOROOT/src/testing.
//
// TODO(adonovan): throws this all away now that x/tools/go/packages
// provides access to the actual synthetic test main files.

import (
	"bytes"
	"fmt"
	"go/ast"
	"go/parser"
	"go/types"
	"log"
	"os"
	"strings"
	"text/template"
)

// FindTests returns the Test, Benchmark, and Example functions
// (as defined by "go test") defined in the specified package,
// and its TestMain function, if any.
//
// Deprecated: Use golang.org/x/tools/go/packages to access synthetic
// testmain packages.
func FindTests(pkg *Package) (tests, benchmarks, examples []*Function, main *Function) {
	prog := pkg.Prog

	// The first two of these may be nil: if the program doesn't import "testing",
	// it can't contain any tests, but it may yet contain Examples.
	var testSig *types.Signature                              // func(*testing.T)
	var benchmarkSig *types.Signature                         // func(*testing.B)
	var exampleSig = types.NewSignature(nil, nil, nil, false) // func()

	// Obtain the types from the parameters of testing.MainStart.
	if testingPkg := prog.ImportedPackage("testing"); testingPkg != nil {
		mainStart := testingPkg.Func("MainStart")
		params := mainStart.Signature.Params()
		testSig = funcField(params.At(1).Type())
		benchmarkSig = funcField(params.At(2).Type())

		// Does the package define this function?
		//   func TestMain(*testing.M)
		if f := pkg.Func("TestMain"); f != nil {
			sig := f.Type().(*types.Signature)
			starM := mainStart.Signature.Results().At(0).Type() // *testing.M
			if sig.Results().Len() == 0 &&
				sig.Params().Len() == 1 &&
				types.Identical(sig.Params().At(0).Type(), starM) {
				main = f
			}
		}
	}

	// TODO(adonovan): use a stable order, e.g. lexical.
	for _, mem := range pkg.Members {
		if f, ok := mem.(*Function); ok &&
			ast.IsExported(f.Name()) &&
			strings.HasSuffix(prog.Fset.Position(f.Pos()).Filename, "_test.go") {

			switch {
			case testSig != nil && isTestSig(f, "Test", testSig):
				tests = append(tests, f)
			case benchmarkSig != nil && isTestSig(f, "Benchmark", benchmarkSig):
				benchmarks = append(benchmarks, f)
			case isTestSig(f, "Example", exampleSig):
				examples = append(examples, f)
			default:
				continue
			}
		}
	}
	return
}

// Like isTest, but checks the signature too.
func isTestSig(f *Function, prefix string, sig *types.Signature) bool {
	return isTest(f.Name(), prefix) && types.Identical(f.Signature, sig)
}

// Given the type of one of the three slice parameters of testing.Main,
// returns the function type.
func funcField(slice types.Type) *types.Signature {
	return slice.(*types.Slice).Elem().Underlying().(*types.Struct).Field(1).Type().(*types.Signature)
}

// isTest tells whether name looks like a test (or benchmark, according to prefix).
// It is a Test (say) if there is a character after Test that is not a lower-case letter.
// We don't want TesticularCancer.
// Plundered from $GOROOT/src/cmd/go/test.go
func isTest(name, prefix string) bool {
	if !strings.HasPrefix(name, prefix) {
		return false
	}
	if len(name) == len(prefix) { // "Test" is ok
		return true
	}
	return ast.IsExported(name[len(prefix):])
}

// CreateTestMainPackage creates and returns a synthetic "testmain"
// package for the specified package if it defines tests, benchmarks or
// executable examples, or nil otherwise.  The new package is named
// "main" and provides a function named "main" that runs the tests,
// similar to the one that would be created by the 'go test' tool.
//
// Subsequent calls to prog.AllPackages include the new package.
// The package pkg must belong to the program prog.
//
// Deprecated: Use golang.org/x/tools/go/packages to access synthetic
// testmain packages.
func (prog *Program) CreateTestMainPackage(pkg *Package) *Package {
	if pkg.Prog != prog {
		log.Fatal("Package does not belong to Program")
	}

	// Template data
	var data struct {
		Pkg                         *Package
		Tests, Benchmarks, Examples []*Function
		Main                        *Function
		Go18                        bool
	}
	data.Pkg = pkg

	// Enumerate tests.
	data.Tests, data.Benchmarks, data.Examples, data.Main = FindTests(pkg)
	if data.Main == nil &&
		data.Tests == nil && data.Benchmarks == nil && data.Examples == nil {
		return nil
	}

	// Synthesize source for testmain package.
	path := pkg.Pkg.Path() + "$testmain"
	tmpl := testmainTmpl
	if testingPkg := prog.ImportedPackage("testing"); testingPkg != nil {
		// In Go 1.8, testing.MainStart's first argument is an interface, not a func.
		data.Go18 = types.IsInterface(testingPkg.Func("MainStart").Signature.Params().At(0).Type())
	} else {
		// The program does not import "testing", but FindTests
		// returned non-nil, which must mean there were Examples
		// but no Test, Benchmark, or TestMain functions.

		// We'll simply call them from testmain.main; this will
		// ensure they don't panic, but will not check any
		// "Output:" comments.
		// (We should not execute an Example that has no
		// "Output:" comment, but it's impossible to tell here.)
		tmpl = examplesOnlyTmpl
	}
	var buf bytes.Buffer
	if err := tmpl.Execute(&buf, data); err != nil {
		log.Fatalf("internal error expanding template for %s: %v", path, err)
	}
	if false { // debugging
		fmt.Fprintln(os.Stderr, buf.String())
	}

	// Parse and type-check the testmain package.
	f, err := parser.ParseFile(prog.Fset, path+".go", &buf, parser.Mode(0))
	if err != nil {
		log.Fatalf("internal error parsing %s: %v", path, err)
	}
	conf := types.Config{
		DisableUnusedImportCheck: true,
		Importer:                 importer{pkg},
	}
	files := []*ast.File{f}
	info := &types.Info{
		Types:      make(map[ast.Expr]types.TypeAndValue),
		Defs:       make(map[*ast.Ident]types.Object),
		Uses:       make(map[*ast.Ident]types.Object),
		Implicits:  make(map[ast.Node]types.Object),
		Scopes:     make(map[ast.Node]*types.Scope),
		Selections: make(map[*ast.SelectorExpr]*types.Selection),
	}
	testmainPkg, err := conf.Check(path, prog.Fset, files, info)
	if err != nil {
		log.Fatalf("internal error type-checking %s: %v", path, err)
	}

	// Create and build SSA code.
	testmain := prog.CreatePackage(testmainPkg, files, info, false)
	testmain.SetDebugMode(false)
	testmain.Build()
	testmain.Func("main").Synthetic = "test main function"
	testmain.Func("init").Synthetic = "package initializer"
	return testmain
}

// An implementation of types.Importer for an already loaded SSA program.
type importer struct {
	pkg *Package // package under test; may be non-importable
}

func (imp importer) Import(path string) (*types.Package, error) {
	if p := imp.pkg.Prog.ImportedPackage(path); p != nil {
		return p.Pkg, nil
	}
	if path == imp.pkg.Pkg.Path() {
		return imp.pkg.Pkg, nil
	}
	return nil, fmt.Errorf("not found") // can't happen
}

var testmainTmpl = template.Must(template.New("testmain").Parse(`
package main

import "io"
import "os"
import "testing"
import p {{printf "%q" .Pkg.Pkg.Path}}

{{if .Go18}}
type deps struct{}

func (deps) ImportPath() string { return "" }
func (deps) MatchString(pat, str string) (bool, error) { return true, nil }
func (deps) StartCPUProfile(io.Writer) error { return nil }
func (deps) StartTestLog(io.Writer) {}
func (deps) StopCPUProfile() {}
func (deps) StopTestLog() error { return nil }
func (deps) WriteHeapProfile(io.Writer) error { return nil }
func (deps) WriteProfileTo(string, io.Writer, int) error { return nil }

var match deps
{{else}}
func match(_, _ string) (bool, error) { return true, nil }
{{end}}

func main() {
	tests := []testing.InternalTest{
{{range .Tests}}
		{ {{printf "%q" .Name}}, p.{{.Name}} },
{{end}}
	}
	benchmarks := []testing.InternalBenchmark{
{{range .Benchmarks}}
		{ {{printf "%q" .Name}}, p.{{.Name}} },
{{end}}
	}
	examples := []testing.InternalExample{
{{range .Examples}}
		{Name: {{printf "%q" .Name}}, F: p.{{.Name}}},
{{end}}
	}
	m := testing.MainStart(match, tests, benchmarks, examples)
{{with .Main}}
	p.{{.Name}}(m)
{{else}}
	os.Exit(m.Run())
{{end}}
}

`))

var examplesOnlyTmpl = template.Must(template.New("examples").Parse(`
package main

import p {{printf "%q" .Pkg.Pkg.Path}}

func main() {
{{range .Examples}}
	p.{{.Name}}()
{{end}}
}
`))
<commit_msg>go/ssa: implement missing testing.testDeps SetPanicOnExit0 method<commit_after>// Copyright 2013 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package ssa

// CreateTestMainPackage synthesizes a main package that runs all the
// tests of the supplied packages.
// It is closely coupled to $GOROOT/src/cmd/go/test.go and $GOROOT/src/testing.
//
// TODO(adonovan): throws this all away now that x/tools/go/packages
// provides access to the actual synthetic test main files.

import (
	"bytes"
	"fmt"
	"go/ast"
	"go/parser"
	"go/types"
	"log"
	"os"
	"strings"
	"text/template"
)

// FindTests returns the Test, Benchmark, and Example functions
// (as defined by "go test") defined in the specified package,
// and its TestMain function, if any.
//
// Deprecated: Use golang.org/x/tools/go/packages to access synthetic
// testmain packages.
func FindTests(pkg *Package) (tests, benchmarks, examples []*Function, main *Function) {
	prog := pkg.Prog

	// The first two of these may be nil: if the program doesn't import "testing",
	// it can't contain any tests, but it may yet contain Examples.
	var testSig *types.Signature                              // func(*testing.T)
	var benchmarkSig *types.Signature                         // func(*testing.B)
	var exampleSig = types.NewSignature(nil, nil, nil, false) // func()

	// Obtain the types from the parameters of testing.MainStart.
	if testingPkg := prog.ImportedPackage("testing"); testingPkg != nil {
		mainStart := testingPkg.Func("MainStart")
		params := mainStart.Signature.Params()
		testSig = funcField(params.At(1).Type())
		benchmarkSig = funcField(params.At(2).Type())

		// Does the package define this function?
		//   func TestMain(*testing.M)
		if f := pkg.Func("TestMain"); f != nil {
			sig := f.Type().(*types.Signature)
			starM := mainStart.Signature.Results().At(0).Type() // *testing.M
			if sig.Results().Len() == 0 &&
				sig.Params().Len() == 1 &&
				types.Identical(sig.Params().At(0).Type(), starM) {
				main = f
			}
		}
	}

	// TODO(adonovan): use a stable order, e.g. lexical.
	for _, mem := range pkg.Members {
		if f, ok := mem.(*Function); ok &&
			ast.IsExported(f.Name()) &&
			strings.HasSuffix(prog.Fset.Position(f.Pos()).Filename, "_test.go") {

			switch {
			case testSig != nil && isTestSig(f, "Test", testSig):
				tests = append(tests, f)
			case benchmarkSig != nil && isTestSig(f, "Benchmark", benchmarkSig):
				benchmarks = append(benchmarks, f)
			case isTestSig(f, "Example", exampleSig):
				examples = append(examples, f)
			default:
				continue
			}
		}
	}
	return
}

// Like isTest, but checks the signature too.
func isTestSig(f *Function, prefix string, sig *types.Signature) bool {
	return isTest(f.Name(), prefix) && types.Identical(f.Signature, sig)
}

// Given the type of one of the three slice parameters of testing.Main,
// returns the function type.
func funcField(slice types.Type) *types.Signature {
	return slice.(*types.Slice).Elem().Underlying().(*types.Struct).Field(1).Type().(*types.Signature)
}

// isTest tells whether name looks like a test (or benchmark, according to prefix).
// It is a Test (say) if there is a character after Test that is not a lower-case letter.
// We don't want TesticularCancer.
// Plundered from $GOROOT/src/cmd/go/test.go
func isTest(name, prefix string) bool {
	if !strings.HasPrefix(name, prefix) {
		return false
	}
	if len(name) == len(prefix) { // "Test" is ok
		return true
	}
	return ast.IsExported(name[len(prefix):])
}

// CreateTestMainPackage creates and returns a synthetic "testmain"
// package for the specified package if it defines tests, benchmarks or
// executable examples, or nil otherwise.  The new package is named
// "main" and provides a function named "main" that runs the tests,
// similar to the one that would be created by the 'go test' tool.
//
// Subsequent calls to prog.AllPackages include the new package.
// The package pkg must belong to the program prog.
//
// Deprecated: Use golang.org/x/tools/go/packages to access synthetic
// testmain packages.
func (prog *Program) CreateTestMainPackage(pkg *Package) *Package {
	if pkg.Prog != prog {
		log.Fatal("Package does not belong to Program")
	}

	// Template data
	var data struct {
		Pkg                         *Package
		Tests, Benchmarks, Examples []*Function
		Main                        *Function
		Go18                        bool
	}
	data.Pkg = pkg

	// Enumerate tests.
	data.Tests, data.Benchmarks, data.Examples, data.Main = FindTests(pkg)
	if data.Main == nil &&
		data.Tests == nil && data.Benchmarks == nil && data.Examples == nil {
		return nil
	}

	// Synthesize source for testmain package.
	path := pkg.Pkg.Path() + "$testmain"
	tmpl := testmainTmpl
	if testingPkg := prog.ImportedPackage("testing"); testingPkg != nil {
		// In Go 1.8, testing.MainStart's first argument is an interface, not a func.
		data.Go18 = types.IsInterface(testingPkg.Func("MainStart").Signature.Params().At(0).Type())
	} else {
		// The program does not import "testing", but FindTests
		// returned non-nil, which must mean there were Examples
		// but no Test, Benchmark, or TestMain functions.

		// We'll simply call them from testmain.main; this will
		// ensure they don't panic, but will not check any
		// "Output:" comments.
		// (We should not execute an Example that has no
		// "Output:" comment, but it's impossible to tell here.)
		tmpl = examplesOnlyTmpl
	}
	var buf bytes.Buffer
	if err := tmpl.Execute(&buf, data); err != nil {
		log.Fatalf("internal error expanding template for %s: %v", path, err)
	}
	if false { // debugging
		fmt.Fprintln(os.Stderr, buf.String())
	}

	// Parse and type-check the testmain package.
	f, err := parser.ParseFile(prog.Fset, path+".go", &buf, parser.Mode(0))
	if err != nil {
		log.Fatalf("internal error parsing %s: %v", path, err)
	}
	conf := types.Config{
		DisableUnusedImportCheck: true,
		Importer:                 importer{pkg},
	}
	files := []*ast.File{f}
	info := &types.Info{
		Types:      make(map[ast.Expr]types.TypeAndValue),
		Defs:       make(map[*ast.Ident]types.Object),
		Uses:       make(map[*ast.Ident]types.Object),
		Implicits:  make(map[ast.Node]types.Object),
		Scopes:     make(map[ast.Node]*types.Scope),
		Selections: make(map[*ast.SelectorExpr]*types.Selection),
	}
	testmainPkg, err := conf.Check(path, prog.Fset, files, info)
	if err != nil {
		log.Fatalf("internal error type-checking %s: %v", path, err)
	}

	// Create and build SSA code.
	testmain := prog.CreatePackage(testmainPkg, files, info, false)
	testmain.SetDebugMode(false)
	testmain.Build()
	testmain.Func("main").Synthetic = "test main function"
	testmain.Func("init").Synthetic = "package initializer"
	return testmain
}

// An implementation of types.Importer for an already loaded SSA program.
type importer struct {
	pkg *Package // package under test; may be non-importable
}

func (imp importer) Import(path string) (*types.Package, error) {
	if p := imp.pkg.Prog.ImportedPackage(path); p != nil {
		return p.Pkg, nil
	}
	if path == imp.pkg.Pkg.Path() {
		return imp.pkg.Pkg, nil
	}
	return nil, fmt.Errorf("not found") // can't happen
}

var testmainTmpl = template.Must(template.New("testmain").Parse(`
package main

import "io"
import "os"
import "testing"
import p {{printf "%q" .Pkg.Pkg.Path}}

{{if .Go18}}
type deps struct{}

func (deps) ImportPath() string { return "" }
func (deps) MatchString(pat, str string) (bool, error) { return true, nil }
func (deps) SetPanicOnExit0(bool) {}
func (deps) StartCPUProfile(io.Writer) error { return nil }
func (deps) StartTestLog(io.Writer) {}
func (deps) StopCPUProfile() {}
func (deps) StopTestLog() error { return nil }
func (deps) WriteHeapProfile(io.Writer) error { return nil }
func (deps) WriteProfileTo(string, io.Writer, int) error { return nil }

var match deps
{{else}}
func match(_, _ string) (bool, error) { return true, nil }
{{end}}

func main() {
	tests := []testing.InternalTest{
{{range .Tests}}
		{ {{printf "%q" .Name}}, p.{{.Name}} },
{{end}}
	}
	benchmarks := []testing.InternalBenchmark{
{{range .Benchmarks}}
		{ {{printf "%q" .Name}}, p.{{.Name}} },
{{end}}
	}
	examples := []testing.InternalExample{
{{range .Examples}}
		{Name: {{printf "%q" .Name}}, F: p.{{.Name}}},
{{end}}
	}
	m := testing.MainStart(match, tests, benchmarks, examples)
{{with .Main}}
	p.{{.Name}}(m)
{{else}}
	os.Exit(m.Run())
{{end}}
}

`))

var examplesOnlyTmpl = template.Must(template.New("examples").Parse(`
package main

import p {{printf "%q" .Pkg.Pkg.Path}}

func main() {
{{range .Examples}}
	p.{{.Name}}()
{{end}}
}
`))
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"os/user"
	"runtime"
	"sort"
	"strings"
	"sync/atomic"
	"time"
)

const (
	helloVersion = "0.4"
)

var knownPaths []string
var boottime time.Time
var banner string
var requests int64
var usr *user.User

func inc() int64 {
	return atomic.AddInt64(&requests, 1)
}

func get() int64 {
	return atomic.LoadInt64(&requests)
}

func getVersion() string {
	return fmt.Sprintf("version=%s runtime=%s pid=%d host=%s GOMAXPROCS=%d", helloVersion, runtime.Version(), os.Getpid(), getHostname(), runtime.GOMAXPROCS(0))
}

func getHostname() string {
	host, errHost := os.Hostname()
	if errHost != nil {
		host = host + " (host error: " + errHost.Error() + ")"
	}
	return host
}

func dumpVersion(path string) {
	log.Printf("dumpVersion: writing to touch=%s", path)
	v := getVersion()
	err := ioutil.WriteFile(path, []byte(v), 0640)
	if err != nil {
		log.Printf("dumpVersion: touch=%s: %v", path, err)
	}
}

func main() {

	boottime = time.Now()

	tls := true

	log.Print(getVersion())

	defaultBanner := "banner default"
	if envBanner := os.Getenv("GWH_BANNER"); envBanner != "" {
		defaultBanner = envBanner
	}

	var errUser error
	usr, errUser = user.Current()
	if errUser != nil {
		log.Printf("current user error: %v", errUser)
	}
	if usr != nil {
		log.Printf("user: %s (uid: %s)", usr.Username, usr.Uid)
	}

	currDir, err := os.Getwd()
	if err != nil {
		log.Fatalf("Getwd: %s", err)
	}

	var addr, httpsAddr, key, cert string
	var disableKeepalive bool
	var touch string
	var showVersion bool

	flag.StringVar(&key, "key", "key.pem", "TLS key file")
	flag.StringVar(&cert, "cert", "cert.pem", "TLS cert file")
	flag.StringVar(&addr, "addr", ":8080", "HTTP listen address")
	flag.StringVar(&httpsAddr, "httpsAddr", ":8443", "HTTPS listen address")
	flag.StringVar(&banner, "banner", defaultBanner, "banner will be displayed (env var GWH_BANNER sets default banner)")
	flag.StringVar(&touch, "touch", "", "write version info into this file, if specified")
	flag.BoolVar(&showVersion, "version", false, "does not actually run")
	flag.BoolVar(&disableKeepalive, "disableKeepalive", false, "disable keepalive")
	flag.Parse()

	if touch != "" {
		dumpVersion(touch)
	}

	keepalive := !disableKeepalive

	log.Print("banner: ", banner)
	log.Print("keepalive: ", keepalive)

	if !fileExists(key) {
		log.Printf("TLS key file not found: %s - disabling TLS", key)
		tls = false
	}

	if !fileExists(cert) {
		log.Printf("TLS cert file not found: %s - disabling TLS", cert)
		tls = false
	}

	if showVersion {
		return
	}

	// default handler
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) { rootHandler(w, r, keepalive) })

	registerStatic("/www/", currDir)

	log.Printf("using TCP ports HTTP=%s HTTPS=%s TLS=%v", addr, httpsAddr, tls)

	if tls {
		serveHTTPS(addr, httpsAddr, cert, key, keepalive)
		return
	}

	log.Printf("serving HTTP on TCP %s", addr)
	if err := listenAndServe(addr, nil, keepalive); err != nil {
		log.Fatalf("listenAndServe: %s: %v", addr, err)
	}
}

func serveHTTPS(addr, httpsAddr, cert, key string, keepalive bool) {
	httpPort := "80"
	h := strings.Split(addr, ":")
	if len(h) > 1 {
		httpPort = h[1]
	}

	httpsPort := "443"
	hs := strings.Split(httpsAddr, ":")
	if len(hs) > 1 {
		httpsPort = hs[1]
	}

	if httpPort != httpsPort {
		log.Printf("installing redirect from HTTP=%s to HTTPS=%s", addr, httpsPort)

		redirectTLS := func(w http.ResponseWriter, r *http.Request) {
			host := strings.Split(r.Host, ":")[0]
			//code := http.StatusMovedPermanently // 301 permanent
			code := http.StatusFound // 302 temporary
			http.Redirect(w, r, "https://"+host+":"+httpsPort+r.RequestURI, code)
		}

		// http-to-https redirect server
		go func() {
			if err := listenAndServe(addr, http.HandlerFunc(redirectTLS), keepalive); err != nil {
				log.Fatalf("redirect: listenAndServe: %s: %v", addr, err)
			}
		}()
	}

	log.Printf("serving HTTPS on TCP %s", httpsAddr)
	if err := listenAndServeTLS(httpsAddr, cert, key, nil, keepalive); err != nil {
		log.Fatalf("listenAndServeTLS: %s: %v", httpsAddr, err)
	}
}

func listenAndServe(addr string, handler http.Handler, keepalive bool) error {
	server := &http.Server{Addr: addr, Handler: handler}
	server.SetKeepAlivesEnabled(keepalive)
	return server.ListenAndServe()
}

func listenAndServeTLS(addr, certFile, keyFile string, handler http.Handler, keepalive bool) error {
	server := &http.Server{Addr: addr, Handler: handler}
	server.SetKeepAlivesEnabled(keepalive)
	return server.ListenAndServeTLS(certFile, keyFile)
}

func fileExists(path string) bool {
	_, err := os.Stat(path)
	return err == nil
}

type staticHandler struct {
	innerHandler http.Handler
}

func registerStatic(path, dir string) {
	http.Handle(path, staticHandler{http.StripPrefix(path, http.FileServer(http.Dir(dir)))})
	knownPaths = append(knownPaths, path)
	log.Printf("mapping www path %s to directory %s", path, dir)
}

func (handler staticHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	count := inc()
	log.Printf("staticHandler.ServeHTTP req=%d url=%s from=%s", count, r.URL.Path, r.RemoteAddr)
	handler.innerHandler.ServeHTTP(w, r)
}

func rootHandler(w http.ResponseWriter, r *http.Request, keepalive bool) {
	count := inc()
	log.Printf("rootHandler: req=%d from=%s", count, r.RemoteAddr)
	log.Printf("rootHandler: req=%d method=%s host=%s path=%s", count, r.Method, r.Host, r.URL.Path)
	log.Printf("rootHandler: req=%d query=%s", count, r.URL.RawQuery)

	var paths string
	for _, p := range knownPaths {
		paths += fmt.Sprintf("%s
", p, p)
	}

	var errMsg string
	if r.URL.Path != "/" {
		format := `

Path not found!

Method: %s

Host: %s

Path not found: [%s]

Query: [%s]

`
		errMsg = fmt.Sprintf(format, r.Method, r.Host, r.URL.Path, r.URL.RawQuery)
	}

	header :=
		`<!DOCTYPE html>

    <title>gowebhello root page</title>

gowebhello root page

    https://github.com/udhos/gowebhello is a simple golang replacement for 'python -m SimpleHTTPServer'.

`
	bodyTempl :=
		`
Welcome!

	gowebhello version %s runtime %s

        Keepalive: %v

	Application banner: %s

	Application arguments: %v

	Application dir: %s

	Process: %d

	User: %s (uid: %s)

	Server hostname: %s

	Your address: %s

	Request method=%s host=%s path=[%s] query=[%s]

	Current time: %s

	Uptime: %s

	Requests: %d

    %s

All known paths:

    %s
`

	footer :=
		`

`

	cwd, errCwd := os.Getwd()
	if errCwd != nil {
		cwd = cwd + " (error: " + errCwd.Error() + ")"
	}

	host := getHostname()

	now := time.Now()

	username := "?"
	uid := "?"

	if usr != nil {
		username = usr.Username
		uid = usr.Uid
	}

	body := fmt.Sprintf(bodyTempl, helloVersion, runtime.Version(), keepalive, banner, os.Args, cwd, os.Getpid(), username, uid, host, r.RemoteAddr, r.Method, r.Host, r.URL.Path, r.URL.RawQuery, now, time.Since(boottime), get(), errMsg, paths)

	if !keepalive {
		w.Header().Set("Connection", "close")
	}

	io.WriteString(w, header)
	io.WriteString(w, body)
	showHeaders(w, r)
	showReqBody(w, r)
	io.WriteString(w, footer)
}

func showReqBody(w http.ResponseWriter, r *http.Request) {

	io.WriteString(w, "
Request Body
\n")

	buf, err := ioutil.ReadAll(r.Body)
	if err != nil {
		m := fmt.Sprintf("request body: %v", err)
		log.Print(m)
		io.WriteString(w, m)
		return
	}

	bodySize := len(buf)

	log.Printf("body size: %d", bodySize)

	io.WriteString(w, fmt.Sprintf("
request body size: %d
\n", bodySize))

	io.WriteString(w, "
\n")
	io.WriteString(w, string(buf))
	io.WriteString(w, "\n")
	io.WriteString(w, "
\n")
}

func showHeaders(w http.ResponseWriter, r *http.Request) {
	io.WriteString(w, "
Headers
\n")

	var headers []string
	for k, v := range r.Header {
		for _, vv := range v {
			headers = append(headers, k+": "+vv+"
\n")
		}
	}

	sort.Strings(headers)

	for _, h := range headers {
		io.WriteString(w, h)
	}
}
<commit_msg>Quota for requests.<commit_after>package main

import (
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"os/user"
	"runtime"
	"sort"
	"strings"
	"sync/atomic"
	"time"
)

const (
	helloVersion = "0.4"
)

var knownPaths []string
var boottime time.Time
var banner string
var requests int64
var quota int64
var usr *user.User

func inc() int64 {
	return atomic.AddInt64(&requests, 1)
}

func get() int64 {
	return atomic.LoadInt64(&requests)
}

func getVersion() string {
	return fmt.Sprintf("version=%s runtime=%s pid=%d host=%s GOMAXPROCS=%d", helloVersion, runtime.Version(), os.Getpid(), getHostname(), runtime.GOMAXPROCS(0))
}

func getHostname() string {
	host, errHost := os.Hostname()
	if errHost != nil {
		host = host + " (host error: " + errHost.Error() + ")"
	}
	return host
}

func dumpVersion(path string) {
	log.Printf("dumpVersion: writing to touch=%s", path)
	v := getVersion()
	err := ioutil.WriteFile(path, []byte(v), 0640)
	if err != nil {
		log.Printf("dumpVersion: touch=%s: %v", path, err)
	}
}

func main() {

	boottime = time.Now()

	tls := true

	log.Print(getVersion())

	defaultBanner := "banner default"
	if envBanner := os.Getenv("GWH_BANNER"); envBanner != "" {
		defaultBanner = envBanner
	}

	var errUser error
	usr, errUser = user.Current()
	if errUser != nil {
		log.Printf("current user error: %v", errUser)
	}
	if usr != nil {
		log.Printf("user: %s (uid: %s)", usr.Username, usr.Uid)
	}

	currDir, err := os.Getwd()
	if err != nil {
		log.Fatalf("Getwd: %s", err)
	}

	var addr, httpsAddr, key, cert string
	var disableKeepalive bool
	var touch string
	var showVersion bool

	flag.StringVar(&key, "key", "key.pem", "TLS key file")
	flag.StringVar(&cert, "cert", "cert.pem", "TLS cert file")
	flag.StringVar(&addr, "addr", ":8080", "HTTP listen address")
	flag.StringVar(&httpsAddr, "httpsAddr", ":8443", "HTTPS listen address")
	flag.StringVar(&banner, "banner", defaultBanner, "banner will be displayed (env var GWH_BANNER sets default banner)")
	flag.StringVar(&touch, "touch", "", "write version info into this file, if specified")
	flag.BoolVar(&showVersion, "version", false, "does not actually run")
	flag.BoolVar(&disableKeepalive, "disableKeepalive", false, "disable keepalive")
	flag.Int64Var(&quota, "quota", 0, "if defined, service is terminated after serving that amount of requests")
	flag.Parse()

	if touch != "" {
		dumpVersion(touch)
	}

	keepalive := !disableKeepalive

	log.Print("banner: ", banner)
	log.Print("keepalive: ", keepalive)

	if !fileExists(key) {
		log.Printf("TLS key file not found: %s - disabling TLS", key)
		tls = false
	}

	if !fileExists(cert) {
		log.Printf("TLS cert file not found: %s - disabling TLS", cert)
		tls = false
	}

	if showVersion {
		return
	}

	// default handler
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) { rootHandler(w, r, keepalive) })

	registerStatic("/www/", currDir)

	log.Printf("using TCP ports HTTP=%s HTTPS=%s TLS=%v", addr, httpsAddr, tls)

	log.Printf("requests quota: %d", quota)

	if tls {
		serveHTTPS(addr, httpsAddr, cert, key, keepalive)
		return
	}

	log.Printf("serving HTTP on TCP %s", addr)
	if err := listenAndServe(addr, nil, keepalive); err != nil {
		log.Fatalf("listenAndServe: %s: %v", addr, err)
	}
}

func serveHTTPS(addr, httpsAddr, cert, key string, keepalive bool) {
	httpPort := "80"
	h := strings.Split(addr, ":")
	if len(h) > 1 {
		httpPort = h[1]
	}

	httpsPort := "443"
	hs := strings.Split(httpsAddr, ":")
	if len(hs) > 1 {
		httpsPort = hs[1]
	}

	if httpPort != httpsPort {
		log.Printf("installing redirect from HTTP=%s to HTTPS=%s", addr, httpsPort)

		redirectTLS := func(w http.ResponseWriter, r *http.Request) {
			host := strings.Split(r.Host, ":")[0]
			//code := http.StatusMovedPermanently // 301 permanent
			code := http.StatusFound // 302 temporary
			http.Redirect(w, r, "https://"+host+":"+httpsPort+r.RequestURI, code)
		}

		// http-to-https redirect server
		go func() {
			if err := listenAndServe(addr, http.HandlerFunc(redirectTLS), keepalive); err != nil {
				log.Fatalf("redirect: listenAndServe: %s: %v", addr, err)
			}
		}()
	}

	log.Printf("serving HTTPS on TCP %s", httpsAddr)
	if err := listenAndServeTLS(httpsAddr, cert, key, nil, keepalive); err != nil {
		log.Fatalf("listenAndServeTLS: %s: %v", httpsAddr, err)
	}
}

func listenAndServe(addr string, handler http.Handler, keepalive bool) error {
	server := &http.Server{Addr: addr, Handler: handler}
	server.SetKeepAlivesEnabled(keepalive)
	return server.ListenAndServe()
}

func listenAndServeTLS(addr, certFile, keyFile string, handler http.Handler, keepalive bool) error {
	server := &http.Server{Addr: addr, Handler: handler}
	server.SetKeepAlivesEnabled(keepalive)
	return server.ListenAndServeTLS(certFile, keyFile)
}

func fileExists(path string) bool {
	_, err := os.Stat(path)
	return err == nil
}

type staticHandler struct {
	innerHandler http.Handler
}

func registerStatic(path, dir string) {
	http.Handle(path, staticHandler{http.StripPrefix(path, http.FileServer(http.Dir(dir)))})
	knownPaths = append(knownPaths, path)
	log.Printf("mapping www path %s to directory %s", path, dir)
}

func (handler staticHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	count := inc()
	log.Printf("staticHandler.ServeHTTP req=%d url=%s from=%s", count, r.URL.Path, r.RemoteAddr)
	handler.innerHandler.ServeHTTP(w, r)
	checkQuota("staticHandler.ServeHTTP", count)
}

func checkQuota(label string, count int64) {
	if count > quota {
		log.Printf("%s req=%d quota=%d EXITING", label, count, quota)
		os.Exit(0)
	}
}

func rootHandler(w http.ResponseWriter, r *http.Request, keepalive bool) {
	count := inc()
	log.Printf("rootHandler: req=%d from=%s", count, r.RemoteAddr)
	log.Printf("rootHandler: req=%d method=%s host=%s path=%s", count, r.Method, r.Host, r.URL.Path)
	log.Printf("rootHandler: req=%d query=%s", count, r.URL.RawQuery)

	var paths string
	for _, p := range knownPaths {
		paths += fmt.Sprintf("%s
", p, p)
	}

	var errMsg string
	if r.URL.Path != "/" {
		format := `

Path not found!

Method: %s

Host: %s

Path not found: [%s]

Query: [%s]

`
		errMsg = fmt.Sprintf(format, r.Method, r.Host, r.URL.Path, r.URL.RawQuery)
	}

	header :=
		`<!DOCTYPE html>

    <title>gowebhello root page</title>

gowebhello root page

    https://github.com/udhos/gowebhello is a simple golang replacement for 'python -m SimpleHTTPServer'.

`
	bodyTempl :=
		`
Welcome!

	gowebhello version %s runtime %s

        Keepalive: %v

	Application banner: %s

	Application arguments: %v

	Application dir: %s

	Process: %d

	User: %s (uid: %s)

	Server hostname: %s

	Your address: %s

	Request method=%s host=%s path=[%s] query=[%s]

	Current time: %s

	Uptime: %s

	Requests: %d

    %s

All known paths:

    %s
`

	footer :=
		`

`

	cwd, errCwd := os.Getwd()
	if errCwd != nil {
		cwd = cwd + " (error: " + errCwd.Error() + ")"
	}

	host := getHostname()

	now := time.Now()

	username := "?"
	uid := "?"

	if usr != nil {
		username = usr.Username
		uid = usr.Uid
	}

	body := fmt.Sprintf(bodyTempl, helloVersion, runtime.Version(), keepalive, banner, os.Args, cwd, os.Getpid(), username, uid, host, r.RemoteAddr, r.Method, r.Host, r.URL.Path, r.URL.RawQuery, now, time.Since(boottime), get(), errMsg, paths)

	if !keepalive {
		w.Header().Set("Connection", "close")
	}

	io.WriteString(w, header)
	io.WriteString(w, body)
	showHeaders(w, r)
	showReqBody(w, r)
	io.WriteString(w, footer)

	checkQuota("rootHandler", count)
}

func showReqBody(w http.ResponseWriter, r *http.Request) {

	io.WriteString(w, "
Request Body
\n")

	buf, err := ioutil.ReadAll(r.Body)
	if err != nil {
		m := fmt.Sprintf("request body: %v", err)
		log.Print(m)
		io.WriteString(w, m)
		return
	}

	bodySize := len(buf)

	log.Printf("body size: %d", bodySize)

	io.WriteString(w, fmt.Sprintf("
request body size: %d
\n", bodySize))

	io.WriteString(w, "
\n")
	io.WriteString(w, string(buf))
	io.WriteString(w, "\n")
	io.WriteString(w, "
\n")
}

func showHeaders(w http.ResponseWriter, r *http.Request) {
	io.WriteString(w, "
Headers
\n")

	var headers []string
	for k, v := range r.Header {
		for _, vv := range v {
			headers = append(headers, k+": "+vv+"
\n")
		}
	}

	sort.Strings(headers)

	for _, h := range headers {
		io.WriteString(w, h)
	}
}
<|endoftext|>