<commit_before>package gorp

import (
	"database/sql/driver"
	"encoding/json"
	"errors"
	"fmt"
	"reflect"
	"strings"
)

// A TypeCaster includes TypeDeffer logic but can also return the SQL
// to cast old types to its new type.
type TypeCaster interface {
	// TypeCast should return a string with zero or one '%s'
	// sequences.  If the string contains a '%s', it will be replaced
	// with the old value; otherwise, the return value will simply be
	// used as the type to cast to in the database.
	TypeCast() string
}

type pgJSON []byte

func (j pgJSON) Value() (driver.Value, error) {
	return []byte(j), nil
}

func (j *pgJSON) Scan(src interface{}) error {
	var source []byte
	switch src.(type) {
	case string:
		source = []byte(src.(string))
	case []byte:
		source = src.([]byte)
	default:
		return errors.New("Incompatible type for pgJSON")
	}
	*j = pgJSON(append((*j)[0:0], source...))
	return nil
}

func (j pgJSON) TypeDef() string {
	return "json"
}

type columnLayout struct {
	FieldName  string `json:"field_name"`
	ColumnName string `json:"column_name"`
	TypeDef    string `json:"type_def"`

	// Values that are only used on the new layout, but are
	// unnecessary for old types.
	gotype   reflect.Type `json:"-"`
	typeCast string       `json:"-"`
}

type tableRecord struct {
	ID          int
	Schemaname  string
	Tablename   string
	Layout      pgJSON
	tableLayout []columnLayout `db:"-"`
}

func (t *tableRecord) TableLayout() []columnLayout {
	if t.tableLayout == nil {
		t.tableLayout = []columnLayout{}
		if err := json.Unmarshal([]byte(t.Layout), &t.tableLayout); err != nil {
			panic(err)
		}
	}
	return t.tableLayout
}

func (t *tableRecord) SetTableLayout(l []columnLayout) {
	t.tableLayout = l
	b, err := json.Marshal(l)
	if err != nil {
		panic(err)
	}
	t.Layout = pgJSON(b)
}

func (t *tableRecord) Merge(l []columnLayout) {
	if t.tableLayout == nil {
		t.SetTableLayout(l)
		return
	}
	for _, newCol := range l {
		shouldAppend := true
		for _, oldCol := range t.tableLayout {
			if newCol.ColumnName == oldCol.ColumnName {
				shouldAppend = false
			}
		}
		if shouldAppend {
			t.tableLayout = append(t.tableLayout, newCol)
		}
	}
}

func ptrToVal(t reflect.Type) reflect.Value {
	if t.Kind() == reflect.Ptr {
		t = t.Elem()
	}
	return reflect.New(t)
}

type MigrationManager struct {
	schemaname string
	tablename  string
	dbMap      *DbMap
	oldTables  []*tableRecord
	newTables  []*tableRecord
}

func (m *MigrationManager) layoutFor(t *TableMap) []columnLayout {
	l := make([]columnLayout, 0, len(t.Columns))
	for _, colMap := range t.Columns {
		if colMap.ColumnName == "-" {
			continue
		}
		var stype string
		orig := ptrToVal(colMap.origtype).Interface()
		dbValue := ptrToVal(colMap.gotype).Interface()
		typer, ok := orig.(TypeDeffer)
		if !ok && colMap.origtype != colMap.gotype {
			typer, ok = dbValue.(TypeDeffer)
		}
		if ok {
			stype = typer.TypeDef()
		} else {
			stype = m.dbMap.Dialect.ToSqlType(colMap.gotype, colMap.MaxSize, colMap.isAutoIncr)
		}
		cast := "%s"
		typeCaster, ok := orig.(TypeCaster)
		if !ok {
			typeCaster, ok = dbValue.(TypeCaster)
		}
		if ok {
			cast = typeCaster.TypeCast()
			if !strings.Contains(cast, "%s") {
				cast = "%s::" + cast
			}
		}

		col := columnLayout{
			FieldName:  colMap.fieldName,
			ColumnName: colMap.ColumnName,
			TypeDef:    stype,
			typeCast:   cast,
			gotype:     colMap.gotype,
		}
		l = append(l, col)
	}
	return l
}

func (m *MigrationManager) addTable(t *TableMap) {
	l := m.layoutFor(t)
	for _, r := range m.newTables {
		if r.Schemaname == t.SchemaName && r.Tablename == t.TableName {
			r.Merge(l)
			l = nil
		}
	}
	if l != nil {
		r := &tableRecord{
			Schemaname: t.SchemaName,
			Tablename:  t.TableName,
		}
		r.SetTableLayout(l)
		m.newTables = append(m.newTables, r)
	}
}

func (m *MigrationManager) newTableRecords() []*tableRecord {
	if m.newTables == nil {
		m.newTables = make([]*tableRecord, 0, len(m.dbMap.tables))
		for _, tableMap := range m.dbMap.tables {
			m.addTable(tableMap)
		}
	}
	return m.newTables
}

func (m *MigrationManager) Migrate() (err error) {
	defer func() {
		if r := recover(); r != nil {
			switch src := r.(type) {
			case error:
				err = src
			default:
				err = fmt.Errorf("Recovered from panic: %v", src)
			}
		}
	}()
	quotedTable := m.dbMap.Dialect.QuotedTableForQuery(m.schemaname, m.tablename)
	_, err = m.dbMap.Select(&m.oldTables, "select * from "+quotedTable)
	if err != nil {
		return err
	}
	return m.run()
}

func (m *MigrationManager) run() error {
	for _, newTable := range m.newTableRecords() {
		found := false
		for _, oldTable := range m.oldTables {
			if oldTable.Schemaname == newTable.Schemaname && oldTable.Tablename == newTable.Tablename {
				found = true
				if err := m.migrateTable(oldTable, newTable); err != nil {
					return err
				}
			}
		}
		if !found {
			if err := m.dbMap.Insert(newTable); err != nil {
				return err
			}
		}
	}
	return nil
}

// layout:
//
// [
//   {
//     "field_name": someName,
//     "column_name": someOtherName,
//     "type_def": someDefinition,
//     "constraints": [],
//   }
// ]
func (m *MigrationManager) migrateTable(oldTable, newTable *tableRecord) (err error) {
	tx, err := m.dbMap.Begin()
	if err != nil {
		panic(err)
	}
	defer func() {
		if err == nil {
			err = tx.Commit()
		} else {
			rollbackErr := tx.Rollback()
			if rollbackErr != nil {
				panic(rollbackErr)
			}
		}
	}()
	if oldTable.Schemaname != newTable.Schemaname || oldTable.Tablename != newTable.Tablename {
		return fmt.Errorf("Unsupported operation: table name change (%s.%s to %s.%s)",
			oldTable.Schemaname,
			oldTable.Tablename,
			newTable.Schemaname,
			newTable.Tablename,
		)
	}
	quotedTable := m.dbMap.Dialect.QuotedTableForQuery(newTable.Schemaname, newTable.Tablename)
	for _, newCol := range newTable.TableLayout() {
		quotedColumn := m.dbMap.Dialect.QuoteField(newCol.ColumnName)
		found := false
		for _, oldCol := range oldTable.TableLayout() {
			if strings.ToLower(oldCol.ColumnName) == strings.ToLower(newCol.ColumnName) {
				found = true
				if oldCol.TypeDef != newCol.TypeDef {
					oldQuotedColumn := m.dbMap.Dialect.QuoteField(newCol.ColumnName + "_type_change_bak")
					sql := "ALTER TABLE " + quotedTable + " RENAME COLUMN " + quotedColumn + " TO " + oldQuotedColumn
					if _, err := tx.Exec(sql); err != nil {
						return err
					}
					sql = "ALTER TABLE " + quotedTable + " ADD COLUMN " + quotedColumn + " " + newCol.TypeDef
					if _, err := tx.Exec(sql); err != nil {
						return err
					}
					sql = "UPDATE " + quotedTable + " SET " + quotedColumn + " = " + fmt.Sprintf(newCol.typeCast, oldQuotedColumn)
					if _, err := tx.Exec(sql); err != nil {
						return err
					}
					sql = "ALTER TABLE " + quotedTable + " DROP COLUMN " + oldQuotedColumn
					if _, err := tx.Exec(sql); err != nil {
						return err
					}
				}
				break
			}
			if oldCol.FieldName == newCol.FieldName {
				found = true
				oldQuotedColumn := m.dbMap.Dialect.QuoteField(oldCol.ColumnName)
				sql := "ALTER TABLE " + quotedTable + " RENAME COLUMN " + oldQuotedColumn + " TO " + quotedColumn
				if _, err := tx.Exec(sql); err != nil {
					return err
				}
				break
			}
		}
		if !found {
			sql := "ALTER TABLE " + quotedTable + " ADD COLUMN " + quotedColumn + " " + newCol.TypeDef
			if _, err := tx.Exec(sql); err != nil {
				return err
			}
			defaultVal := reflect.New(newCol.gotype).Interface()
			sql = "UPDATE " + quotedTable + " SET " + quotedColumn + "=" + m.dbMap.Dialect.BindVar(0)
			if _, err := tx.Exec(sql, defaultVal); err != nil {
				return err
			}
		}
	}
	newTable.ID = oldTable.ID
	if count, err := tx.Update(newTable); err != nil || count != 1 {
		return err
	}
	return nil
}

// Migrater returns a MigrationManager using the given tablename as
// the migration table.
func (m *DbMap) Migrater(schemaname, tablename string) (*MigrationManager, error) {
	// Just run the create table statement for the migration table,
	// using a temporary DbMap
	tmpM := &DbMap{
		Db:      m.Db,
		Dialect: m.Dialect,
	}
	tmpM.AddTableWithNameAndSchema(tableRecord{}, schemaname, tablename).SetKeys(true, "ID")
	if err := tmpM.CreateTablesIfNotExists(); err != nil {
		return nil, err
	}
	m.AddTableWithNameAndSchema(tableRecord{}, schemaname, tablename).SetKeys(true, "ID")
	return &MigrationManager{
		schemaname: schemaname,
		tablename:  tablename,
		dbMap:      m,
	}, nil
}
<commit_msg>Add TypeDefSwitcher which is a TypeDeffer that wants to switch its type on the first migration run<commit_after>package gorp

import (
	"database/sql/driver"
	"encoding/json"
	"errors"
	"fmt"
	"reflect"
	"strings"
)

// A TypeDefSwitcher is a TypeDeffer except that it won't have its type
// defined on table creation.  Instead, its type will be changed on
// the first migration run.  Useful for foreign keys when the target
// table may be created later than the column.
type TypeDefSwitcher interface {
	// TypeDefSwitch should return the same thing as
	// TypeDeffer.TypeDef.
	TypeDefSwitch() string
}

// A TypeCaster includes TypeDeffer logic but can also return the SQL
// to cast old types to its new type.
type TypeCaster interface {
	// TypeCast should return a string with zero or one '%s'
	// sequences.  If the string contains a '%s', it will be replaced
	// with the old value; otherwise, the return value will simply be
	// used as the type to cast to in the database.
	TypeCast() string
}

type pgJSON []byte

func (j pgJSON) Value() (driver.Value, error) {
	return []byte(j), nil
}

func (j *pgJSON) Scan(src interface{}) error {
	var source []byte
	switch src.(type) {
	case string:
		source = []byte(src.(string))
	case []byte:
		source = src.([]byte)
	default:
		return errors.New("Incompatible type for pgJSON")
	}
	*j = pgJSON(append((*j)[0:0], source...))
	return nil
}

func (j pgJSON) TypeDef() string {
	return "json"
}

type columnLayout struct {
	FieldName  string `json:"field_name"`
	ColumnName string `json:"column_name"`
	TypeDef    string `json:"type_def"`

	// Values that are only used on the new layout, but are
	// unnecessary for old types.
	isTypeSwitch bool         `json:"-"`
	gotype       reflect.Type `json:"-"`
	typeCast     string       `json:"-"`
}

type tableRecord struct {
	ID          int
	Schemaname  string
	Tablename   string
	Layout      pgJSON
	tableLayout []columnLayout `db:"-"`
}

func (t *tableRecord) TableLayout() []columnLayout {
	if t.tableLayout == nil {
		t.tableLayout = []columnLayout{}
		if err := json.Unmarshal([]byte(t.Layout), &t.tableLayout); err != nil {
			panic(err)
		}
	}
	return t.tableLayout
}

func (t *tableRecord) SetTableLayout(l []columnLayout) {
	t.tableLayout = l
	b, err := json.Marshal(l)
	if err != nil {
		panic(err)
	}
	t.Layout = pgJSON(b)
}

func (t *tableRecord) Merge(l []columnLayout) {
	if t.tableLayout == nil {
		t.SetTableLayout(l)
		return
	}
	for _, newCol := range l {
		shouldAppend := true
		for _, oldCol := range t.tableLayout {
			if newCol.ColumnName == oldCol.ColumnName {
				shouldAppend = false
			}
		}
		if shouldAppend {
			t.tableLayout = append(t.tableLayout, newCol)
		}
	}
}

func ptrToVal(t reflect.Type) reflect.Value {
	if t.Kind() == reflect.Ptr {
		t = t.Elem()
	}
	return reflect.New(t)
}

type MigrationManager struct {
	schemaname string
	tablename  string
	dbMap      *DbMap
	oldTables  []*tableRecord
	newTables  []*tableRecord
}

func (m *MigrationManager) layoutFor(t *TableMap) []columnLayout {
	l := make([]columnLayout, 0, len(t.Columns))
	for _, colMap := range t.Columns {
		if colMap.ColumnName == "-" {
			continue
		}
		var stype string
		orig := ptrToVal(colMap.origtype).Interface()
		dbValue := ptrToVal(colMap.gotype).Interface()
		typer, hasDef := orig.(TypeDeffer)
		if !hasDef && colMap.origtype != colMap.gotype {
			typer, hasDef = dbValue.(TypeDeffer)
		}
		typeSwitcher, hasSwitch := orig.(TypeDefSwitcher)
		if !hasSwitch && colMap.origtype != colMap.gotype {
			typeSwitcher, hasSwitch = dbValue.(TypeDefSwitcher)
		}
		if hasDef {
			stype = typer.TypeDef()
		} else if hasSwitch {
			stype = typeSwitcher.TypeDefSwitch()
		} else {
			stype = m.dbMap.Dialect.ToSqlType(colMap.gotype, colMap.MaxSize, colMap.isAutoIncr)
		}
		cast := "%s"
		typeCaster, ok := orig.(TypeCaster)
		if !ok {
			typeCaster, ok = dbValue.(TypeCaster)
		}
		if ok {
			cast = typeCaster.TypeCast()
			if !strings.Contains(cast, "%s") {
				cast = "%s::" + cast
			}
		}

		col := columnLayout{
			FieldName:    colMap.fieldName,
			ColumnName:   colMap.ColumnName,
			TypeDef:      stype,
			isTypeSwitch: hasSwitch,
			typeCast:     cast,
			gotype:       colMap.gotype,
		}
		l = append(l, col)
	}
	return l
}

func (m *MigrationManager) addTable(t *TableMap) {
	l := m.layoutFor(t)
	for _, r := range m.newTables {
		if r.Schemaname == t.SchemaName && r.Tablename == t.TableName {
			r.Merge(l)
			return
		}
	}
	r := &tableRecord{
		Schemaname: t.SchemaName,
		Tablename:  t.TableName,
	}
	r.SetTableLayout(l)
	m.newTables = append(m.newTables, r)
}

func (m *MigrationManager) newTableRecords() []*tableRecord {
	if m.newTables == nil {
		m.newTables = make([]*tableRecord, 0, len(m.dbMap.tables))
		for _, tableMap := range m.dbMap.tables {
			m.addTable(tableMap)
		}
	}
	return m.newTables
}

func (m *MigrationManager) Migrate() (err error) {
	defer func() {
		if r := recover(); r != nil {
			switch src := r.(type) {
			case error:
				err = src
			default:
				err = fmt.Errorf("Recovered from panic: %v", src)
			}
		}
	}()
	quotedTable := m.dbMap.Dialect.QuotedTableForQuery(m.schemaname, m.tablename)
	_, err = m.dbMap.Select(&m.oldTables, "select * from "+quotedTable)
	if err != nil {
		return err
	}
	return m.run()
}

func (m *MigrationManager) run() error {
	for _, newTable := range m.newTableRecords() {
		found := false
		for _, oldTable := range m.oldTables {
			if oldTable.Schemaname == newTable.Schemaname && oldTable.Tablename == newTable.Tablename {
				found = true
				if err := m.migrateTable(oldTable, newTable); err != nil {
					return err
				}
			}
		}
		if !found {
			m.handleTypeSwitches(newTable)
			if err := m.dbMap.Insert(newTable); err != nil {
				return err
			}
		}
	}
	return nil
}

func (m *MigrationManager) handleTypeSwitches(table *tableRecord) (err error) {
	tx, err := m.dbMap.Begin()
	if err != nil {
		panic(err)
	}
	defer func() {
		if err == nil {
			err = tx.Commit()
		} else {
			if rollbackErr := tx.Rollback(); rollbackErr != nil {
				panic(rollbackErr)
			}
		}
	}()
	quotedTable := m.dbMap.Dialect.QuotedTableForQuery(table.Schemaname, table.Tablename)
	for _, newCol := range table.TableLayout() {
		if newCol.isTypeSwitch {
			if err = m.changeType(quotedTable, newCol, tx); err != nil {
				return err
			}
		}
	}
	return nil
}

func (m *MigrationManager) changeType(quotedTable string, newCol columnLayout, tx *Transaction) error {
	quotedColumn := m.dbMap.Dialect.QuoteField(newCol.ColumnName)
	oldQuotedColumn := m.dbMap.Dialect.QuoteField(newCol.ColumnName + "_type_change_bak")
	sql := "ALTER TABLE " + quotedTable + " RENAME COLUMN " + quotedColumn + " TO " + oldQuotedColumn
	if _, err := tx.Exec(sql); err != nil {
		return err
	}
	sql = "ALTER TABLE " + quotedTable + " ADD COLUMN " + quotedColumn + " " + newCol.TypeDef
	if _, err := tx.Exec(sql); err != nil {
		return err
	}
	sql = "UPDATE " + quotedTable + " SET " + quotedColumn + " = " + fmt.Sprintf(newCol.typeCast, oldQuotedColumn)
	if _, err := tx.Exec(sql); err != nil {
		return err
	}
	sql = "ALTER TABLE " + quotedTable + " DROP COLUMN " + oldQuotedColumn
	if _, err := tx.Exec(sql); err != nil {
		return err
	}
	return nil
}

func (m *MigrationManager) migrateTable(oldTable, newTable *tableRecord) (err error) {
	tx, err := m.dbMap.Begin()
	if err != nil {
		panic(err)
	}
	defer func() {
		if err == nil {
			err = tx.Commit()
		} else {
			if rollbackErr := tx.Rollback(); rollbackErr != nil {
				panic(rollbackErr)
			}
		}
	}()
	if oldTable.Schemaname != newTable.Schemaname || oldTable.Tablename != newTable.Tablename {
		return fmt.Errorf("Unsupported operation: table name change (%s.%s to %s.%s)",
			oldTable.Schemaname,
			oldTable.Tablename,
			newTable.Schemaname,
			newTable.Tablename,
		)
	}
	quotedTable := m.dbMap.Dialect.QuotedTableForQuery(newTable.Schemaname, newTable.Tablename)
	for _, newCol := range newTable.TableLayout() {
		quotedColumn := m.dbMap.Dialect.QuoteField(newCol.ColumnName)
		found := false
		for _, oldCol := range oldTable.TableLayout() {
			if strings.ToLower(oldCol.ColumnName) == strings.ToLower(newCol.ColumnName) {
				found = true
				if oldCol.TypeDef != newCol.TypeDef {
					if err := m.changeType(quotedTable, newCol, tx); err != nil {
						return err
					}
				}
				break
			}
			if oldCol.FieldName == newCol.FieldName {
				found = true
				oldQuotedColumn := m.dbMap.Dialect.QuoteField(oldCol.ColumnName)
				sql := "ALTER TABLE " + quotedTable + " RENAME COLUMN " + oldQuotedColumn + " TO " + quotedColumn
				if _, err := tx.Exec(sql); err != nil {
					return err
				}
				break
			}
		}
		if !found {
			sql := "ALTER TABLE " + quotedTable + " ADD COLUMN " + quotedColumn + " " + newCol.TypeDef
			if _, err := tx.Exec(sql); err != nil {
				return err
			}
			defaultVal := reflect.New(newCol.gotype).Interface()
			sql = "UPDATE " + quotedTable + " SET " + quotedColumn + "=" + m.dbMap.Dialect.BindVar(0)
			if _, err := tx.Exec(sql, defaultVal); err != nil {
				return err
			}
		}
	}
	newTable.ID = oldTable.ID
	if count, err := tx.Update(newTable); err != nil || count != 1 {
		return err
	}
	return nil
}

// Migrater returns a MigrationManager using the given tablename as
// the migration table.
func (m *DbMap) Migrater(schemaname, tablename string) (*MigrationManager, error) {
	// Just run the create table statement for the migration table,
	// using a temporary DbMap
	tmpM := &DbMap{
		Db:      m.Db,
		Dialect: m.Dialect,
	}
	tmpM.AddTableWithNameAndSchema(tableRecord{}, schemaname, tablename).SetKeys(true, "ID")
	if err := tmpM.CreateTablesIfNotExists(); err != nil {
		return nil, err
	}
	m.AddTableWithNameAndSchema(tableRecord{}, schemaname, tablename).SetKeys(true, "ID")
	return &MigrationManager{
		schemaname: schemaname,
		tablename:  tablename,
		dbMap:      m,
	}, nil
}
<|endoftext|>


<commit_before>// Package confirm implements confirmation of user registration via e-mail
package confirm

import (
	"context"
	"crypto/rand"
	"crypto/sha512"
	"crypto/subtle"
	"encoding/base64"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"path"

	"github.com/pkg/errors"

	"github.com/volatiletech/authboss"
)

const (
	// PageConfirm is only really used for the BodyReader
	PageConfirm = "confirm"

	// EmailConfirmHTML is the name of the html template for e-mails
	EmailConfirmHTML = "confirm_html"
	// EmailConfirmTxt is the name of the text template for e-mails
	EmailConfirmTxt = "confirm_txt"

	// FormValueConfirm is the name of the form value for
	FormValueConfirm = "cnf"

	// DataConfirmURL is the name of the e-mail template variable
	// that gives the url to send to the user for confirmation.
	DataConfirmURL = "url"

	confirmTokenSize  = 64
	confirmTokenSplit = confirmTokenSize / 2
)

func init() {
	authboss.RegisterModule("confirm", &Confirm{})
}

// Confirm module
type Confirm struct {
	*authboss.Authboss
}

// Init module
func (c *Confirm) Init(ab *authboss.Authboss) (err error) {
	c.Authboss = ab

	if err = c.Authboss.Config.Core.MailRenderer.Load(EmailConfirmHTML, EmailConfirmTxt); err != nil {
		return err
	}

	var callbackMethod func(string, http.Handler)
	methodConfig := c.Config.Modules.ConfirmMethod
	if methodConfig == http.MethodGet {
		methodConfig = c.Config.Modules.MailRouteMethod
	}
	switch methodConfig {
	case http.MethodGet:
		callbackMethod = c.Authboss.Config.Core.Router.Get
	case http.MethodPost:
		callbackMethod = c.Authboss.Config.Core.Router.Post
	}
	callbackMethod("/confirm", c.Authboss.Config.Core.ErrorHandler.Wrap(c.Get))

	c.Events.Before(authboss.EventAuth, c.PreventAuth)
	c.Events.After(authboss.EventRegister, c.StartConfirmationWeb)

	return nil
}

// PreventAuth stops the EventAuth from succeeding when a user is not confirmed
// This relies on the fact that the context holds the user at this point in time
// loaded by the auth module (or something else).
func (c *Confirm) PreventAuth(w http.ResponseWriter, r *http.Request, handled bool) (bool, error) {
	logger := c.Authboss.RequestLogger(r)

	user, err := c.Authboss.CurrentUser(r)
	if err != nil {
		return false, err
	}

	cuser := authboss.MustBeConfirmable(user)
	if cuser.GetConfirmed() {
		logger.Infof("user %s is confirmed, allowing auth", user.GetPID())
		return false, nil
	}

	logger.Infof("user %s was not confirmed, preventing auth", user.GetPID())
	ro := authboss.RedirectOptions{
		Code:         http.StatusTemporaryRedirect,
		RedirectPath: c.Authboss.Config.Paths.ConfirmNotOK,
		Failure:      "Your account has not been confirmed, please check your e-mail.",
	}
	return true, c.Authboss.Config.Core.Redirector.Redirect(w, r, ro)
}

// StartConfirmationWeb hijacks a request and forces a user to be confirmed
// first it's assumed that the current user is loaded into the request context.
func (c *Confirm) StartConfirmationWeb(w http.ResponseWriter, r *http.Request, handled bool) (bool, error) {
	user, err := c.Authboss.CurrentUser(r)
	if err != nil {
		return false, err
	}

	cuser := authboss.MustBeConfirmable(user)
	if err = c.StartConfirmation(r.Context(), cuser, true); err != nil {
		return false, err
	}

	ro := authboss.RedirectOptions{
		Code:         http.StatusTemporaryRedirect,
		RedirectPath: c.Authboss.Config.Paths.ConfirmNotOK,
		Success:      "Please verify your account, an e-mail has been sent to you.",
	}
	return true, c.Authboss.Config.Core.Redirector.Redirect(w, r, ro)
}

// StartConfirmation begins confirmation on a user by setting them to require
// confirmation via a created token, and optionally sending them an e-mail.
func (c *Confirm) StartConfirmation(ctx context.Context, user authboss.ConfirmableUser, sendEmail bool) error {
	logger := c.Authboss.Logger(ctx)

	selector, verifier, token, err := GenerateConfirmCreds()
	if err != nil {
		return err
	}

	user.PutConfirmed(false)
	user.PutConfirmSelector(selector)
	user.PutConfirmVerifier(verifier)

	logger.Infof("generated new confirm token for user: %s", user.GetPID())
	if err := c.Authboss.Config.Storage.Server.Save(ctx, user); err != nil {
		return errors.Wrap(err, "failed to save user during StartConfirmation, user data may be in weird state")
	}

	goConfirmEmail(c, ctx, user.GetEmail(), token)

	return nil
}

// This is here so it can be mocked out by a test
var goConfirmEmail = func(c *Confirm, ctx context.Context, to, token string) {
	go c.SendConfirmEmail(ctx, to, token)
}

// SendConfirmEmail sends a confirmation e-mail to a user
func (c *Confirm) SendConfirmEmail(ctx context.Context, to, token string) {
	logger := c.Authboss.Logger(ctx)

	mailURL := c.mailURL(token)

	email := authboss.Email{
		To:       []string{to},
		From:     c.Config.Mail.From,
		FromName: c.Config.Mail.FromName,
		Subject:  c.Config.Mail.SubjectPrefix + "Confirm New Account",
	}

	logger.Infof("sending confirm e-mail to: %s", to)

	ro := authboss.EmailResponseOptions{
		Data:         authboss.NewHTMLData(DataConfirmURL, mailURL),
		HTMLTemplate: EmailConfirmHTML,
		TextTemplate: EmailConfirmTxt,
	}
	if err := c.Authboss.Email(ctx, email, ro); err != nil {
		logger.Errorf("failed to send confirm e-mail to %s: %+v", to, err)
	}
}

// Get is a request that confirms a user with a valid token
func (c *Confirm) Get(w http.ResponseWriter, r *http.Request) error {
	logger := c.RequestLogger(r)

	validator, err := c.Authboss.Config.Core.BodyReader.Read(PageConfirm, r)
	if err != nil {
		return err
	}

	if errs := validator.Validate(); errs != nil {
		logger.Infof("validation failed in Confirm.Get, this typically means a bad token: %+v", errs)
		return c.invalidToken(w, r)
	}

	values := authboss.MustHaveConfirmValues(validator)

	rawToken, err := base64.URLEncoding.DecodeString(values.GetToken())
	if err != nil {
		logger.Infof("error decoding token in Confirm.Get, this typically means a bad token: %s %+v", values.GetToken(), err)
		return c.invalidToken(w, r)
	}

	if len(rawToken) != confirmTokenSize {
		logger.Infof("invalid confirm token submitted, size was wrong: %d", len(rawToken))
		return c.invalidToken(w, r)
	}

	selectorBytes := sha512.Sum512(rawToken[:confirmTokenSplit])
	verifierBytes := sha512.Sum512(rawToken[confirmTokenSplit:])
	selector := base64.StdEncoding.EncodeToString(selectorBytes[:])

	storer := authboss.EnsureCanConfirm(c.Authboss.Config.Storage.Server)
	user, err := storer.LoadByConfirmSelector(r.Context(), selector)
	if err == authboss.ErrUserNotFound {
		logger.Infof("confirm selector was not found in database: %s", selector)
		return c.invalidToken(w, r)
	} else if err != nil {
		return err
	}

	dbVerifierBytes, err := base64.StdEncoding.DecodeString(user.GetConfirmVerifier())
	if err != nil {
		logger.Infof("invalid confirm verifier stored in database: %s", user.GetConfirmVerifier())
		return c.invalidToken(w, r)
	}

	if subtle.ConstantTimeEq(int32(len(verifierBytes)), int32(len(dbVerifierBytes))) != 1 ||
		subtle.ConstantTimeCompare(verifierBytes[:], dbVerifierBytes) != 1 {
		logger.Info("stored confirm verifier does not match provided one")
		return c.invalidToken(w, r)
	}

	user.PutConfirmSelector("")
	user.PutConfirmVerifier("")
	user.PutConfirmed(true)

	logger.Infof("user %s confirmed their account", user.GetPID())
	if err = c.Authboss.Config.Storage.Server.Save(r.Context(), user); err != nil {
		return err
	}

	ro := authboss.RedirectOptions{
		Code:         http.StatusTemporaryRedirect,
		Success:      "You have successfully confirmed your account.",
		RedirectPath: c.Authboss.Config.Paths.ConfirmOK,
	}
	return c.Authboss.Config.Core.Redirector.Redirect(w, r, ro)
}

func (c *Confirm) mailURL(token string) string {
	query := url.Values{FormValueConfirm: []string{token}}

	if len(c.Config.Mail.RootURL) != 0 {
		return fmt.Sprintf("%s?%s", c.Config.Mail.RootURL+"/confirm", query.Encode())
	}

	p := path.Join(c.Config.Paths.Mount, "confirm")
	return fmt.Sprintf("%s%s?%s", c.Config.Paths.RootURL, p, query.Encode())
}

func (c *Confirm) invalidToken(w http.ResponseWriter, r *http.Request) error {
	ro := authboss.RedirectOptions{
		Code:         http.StatusTemporaryRedirect,
		Failure:      "confirm token is invalid",
		RedirectPath: c.Authboss.Config.Paths.ConfirmNotOK,
	}
	return c.Authboss.Config.Core.Redirector.Redirect(w, r, ro)
}

// Middleware ensures that a user is confirmed, or else it will intercept the
// request and send them to the confirm page, this will load the user if he's
// not been loaded yet from the session.
//
// Panics if the user was not able to be loaded in order to allow a panic
// handler to show a nice error page, also panics if it failed to redirect
// for whatever reason.
func Middleware(ab *authboss.Authboss) func(http.Handler) http.Handler {
	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			user := ab.LoadCurrentUserP(&r)

			cu := authboss.MustBeConfirmable(user)
			if cu.GetConfirmed() {
				next.ServeHTTP(w, r)
				return
			}

			logger := ab.RequestLogger(r)
			logger.Infof("user %s prevented from accessing %s: not confirmed", user.GetPID(), r.URL.Path)
			ro := authboss.RedirectOptions{
				Code:         http.StatusTemporaryRedirect,
				Failure:      "Your account has not been confirmed, please check your e-mail.",
				RedirectPath: ab.Config.Paths.ConfirmNotOK,
			}
			if err := ab.Config.Core.Redirector.Redirect(w, r, ro); err != nil {
				logger.Errorf("error redirecting in confirm.Middleware: #%v", err)
			}
		})
	}
}

// GenerateConfirmCreds generates pieces needed for user confirm
// selector: hash of the first half of a 64 byte value
// (to be stored in the database and used in SELECT query)
// verifier: hash of the second half of a 64 byte value
// (to be stored in database but never used in SELECT query)
// token: the user-facing base64 encoded selector+verifier
func GenerateConfirmCreds() (selector, verifier, token string, err error) {
	rawToken := make([]byte, confirmTokenSize)
	if _, err = io.ReadFull(rand.Reader, rawToken); err != nil {
		return "", "", "", err
	}
	selectorBytes := sha512.Sum512(rawToken[:confirmTokenSplit])
	verifierBytes := sha512.Sum512(rawToken[confirmTokenSplit:])

	return base64.StdEncoding.EncodeToString(selectorBytes[:]),
		base64.StdEncoding.EncodeToString(verifierBytes[:]),
		base64.URLEncoding.EncodeToString(rawToken),
		nil
}
<commit_msg>Fix potential misconfiguration<commit_after>// Package confirm implements confirmation of user registration via e-mail
package confirm

import (
	"context"
	"crypto/rand"
	"crypto/sha512"
	"crypto/subtle"
	"encoding/base64"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"path"

	"github.com/pkg/errors"

	"github.com/volatiletech/authboss"
)

const (
	// PageConfirm is only really used for the BodyReader
	PageConfirm = "confirm"

	// EmailConfirmHTML is the name of the html template for e-mails
	EmailConfirmHTML = "confirm_html"
	// EmailConfirmTxt is the name of the text template for e-mails
	EmailConfirmTxt = "confirm_txt"

	// FormValueConfirm is the name of the form value for
	FormValueConfirm = "cnf"

	// DataConfirmURL is the name of the e-mail template variable
	// that gives the url to send to the user for confirmation.
	DataConfirmURL = "url"

	confirmTokenSize  = 64
	confirmTokenSplit = confirmTokenSize / 2
)

func init() {
	authboss.RegisterModule("confirm", &Confirm{})
}

// Confirm module
type Confirm struct {
	*authboss.Authboss
}

// Init module
func (c *Confirm) Init(ab *authboss.Authboss) (err error) {
	c.Authboss = ab

	if err = c.Authboss.Config.Core.MailRenderer.Load(EmailConfirmHTML, EmailConfirmTxt); err != nil {
		return err
	}

	var callbackMethod func(string, http.Handler)
	methodConfig := c.Config.Modules.ConfirmMethod
	if methodConfig == http.MethodGet {
		methodConfig = c.Config.Modules.MailRouteMethod
	}
	switch methodConfig {
	case http.MethodGet:
		callbackMethod = c.Authboss.Config.Core.Router.Get
	case http.MethodPost:
		callbackMethod = c.Authboss.Config.Core.Router.Post
	default:
		panic("invalid config for ConfirmMethod/MailRouteMethod")
	}
	callbackMethod("/confirm", c.Authboss.Config.Core.ErrorHandler.Wrap(c.Get))

	c.Events.Before(authboss.EventAuth, c.PreventAuth)
	c.Events.After(authboss.EventRegister, c.StartConfirmationWeb)

	return nil
}

// PreventAuth stops the EventAuth from succeeding when a user is not confirmed
// This relies on the fact that the context holds the user at this point in time
// loaded by the auth module (or something else).
func (c *Confirm) PreventAuth(w http.ResponseWriter, r *http.Request, handled bool) (bool, error) {
	logger := c.Authboss.RequestLogger(r)

	user, err := c.Authboss.CurrentUser(r)
	if err != nil {
		return false, err
	}

	cuser := authboss.MustBeConfirmable(user)
	if cuser.GetConfirmed() {
		logger.Infof("user %s is confirmed, allowing auth", user.GetPID())
		return false, nil
	}

	logger.Infof("user %s was not confirmed, preventing auth", user.GetPID())
	ro := authboss.RedirectOptions{
		Code:         http.StatusTemporaryRedirect,
		RedirectPath: c.Authboss.Config.Paths.ConfirmNotOK,
		Failure:      "Your account has not been confirmed, please check your e-mail.",
	}
	return true, c.Authboss.Config.Core.Redirector.Redirect(w, r, ro)
}

// StartConfirmationWeb hijacks a request and forces a user to be confirmed
// first it's assumed that the current user is loaded into the request context.
func (c *Confirm) StartConfirmationWeb(w http.ResponseWriter, r *http.Request, handled bool) (bool, error) {
	user, err := c.Authboss.CurrentUser(r)
	if err != nil {
		return false, err
	}

	cuser := authboss.MustBeConfirmable(user)
	if err = c.StartConfirmation(r.Context(), cuser, true); err != nil {
		return false, err
	}

	ro := authboss.RedirectOptions{
		Code:         http.StatusTemporaryRedirect,
		RedirectPath: c.Authboss.Config.Paths.ConfirmNotOK,
		Success:      "Please verify your account, an e-mail has been sent to you.",
	}
	return true, c.Authboss.Config.Core.Redirector.Redirect(w, r, ro)
}

// StartConfirmation begins confirmation on a user by setting them to require
// confirmation via a created token, and optionally sending them an e-mail.
func (c *Confirm) StartConfirmation(ctx context.Context, user authboss.ConfirmableUser, sendEmail bool) error {
	logger := c.Authboss.Logger(ctx)

	selector, verifier, token, err := GenerateConfirmCreds()
	if err != nil {
		return err
	}

	user.PutConfirmed(false)
	user.PutConfirmSelector(selector)
	user.PutConfirmVerifier(verifier)

	logger.Infof("generated new confirm token for user: %s", user.GetPID())
	if err := c.Authboss.Config.Storage.Server.Save(ctx, user); err != nil {
		return errors.Wrap(err, "failed to save user during StartConfirmation, user data may be in weird state")
	}

	goConfirmEmail(c, ctx, user.GetEmail(), token)

	return nil
}

// This is here so it can be mocked out by a test
var goConfirmEmail = func(c *Confirm, ctx context.Context, to, token string) {
	go c.SendConfirmEmail(ctx, to, token)
}

// SendConfirmEmail sends a confirmation e-mail to a user
func (c *Confirm) SendConfirmEmail(ctx context.Context, to, token string) {
	logger := c.Authboss.Logger(ctx)

	mailURL := c.mailURL(token)

	email := authboss.Email{
		To:       []string{to},
		From:     c.Config.Mail.From,
		FromName: c.Config.Mail.FromName,
		Subject:  c.Config.Mail.SubjectPrefix + "Confirm New Account",
	}

	logger.Infof("sending confirm e-mail to: %s", to)

	ro := authboss.EmailResponseOptions{
		Data:         authboss.NewHTMLData(DataConfirmURL, mailURL),
		HTMLTemplate: EmailConfirmHTML,
		TextTemplate: EmailConfirmTxt,
	}
	if err := c.Authboss.Email(ctx, email, ro); err != nil {
		logger.Errorf("failed to send confirm e-mail to %s: %+v", to, err)
	}
}

// Get is a request that confirms a user with a valid token
func (c *Confirm) Get(w http.ResponseWriter, r *http.Request) error {
	logger := c.RequestLogger(r)

	validator, err := c.Authboss.Config.Core.BodyReader.Read(PageConfirm, r)
	if err != nil {
		return err
	}

	if errs := validator.Validate(); errs != nil {
		logger.Infof("validation failed in Confirm.Get, this typically means a bad token: %+v", errs)
		return c.invalidToken(w, r)
	}

	values := authboss.MustHaveConfirmValues(validator)

	rawToken, err := base64.URLEncoding.DecodeString(values.GetToken())
	if err != nil {
		logger.Infof("error decoding token in Confirm.Get, this typically means a bad token: %s %+v", values.GetToken(), err)
		return c.invalidToken(w, r)
	}

	if len(rawToken) != confirmTokenSize {
		logger.Infof("invalid confirm token submitted, size was wrong: %d", len(rawToken))
		return c.invalidToken(w, r)
	}

	selectorBytes := sha512.Sum512(rawToken[:confirmTokenSplit])
	verifierBytes := sha512.Sum512(rawToken[confirmTokenSplit:])
	selector := base64.StdEncoding.EncodeToString(selectorBytes[:])

	storer := authboss.EnsureCanConfirm(c.Authboss.Config.Storage.Server)
	user, err := storer.LoadByConfirmSelector(r.Context(), selector)
	if err == authboss.ErrUserNotFound {
		logger.Infof("confirm selector was not found in database: %s", selector)
		return c.invalidToken(w, r)
	} else if err != nil {
		return err
	}

	dbVerifierBytes, err := base64.StdEncoding.DecodeString(user.GetConfirmVerifier())
	if err != nil {
		logger.Infof("invalid confirm verifier stored in database: %s", user.GetConfirmVerifier())
		return c.invalidToken(w, r)
	}

	if subtle.ConstantTimeEq(int32(len(verifierBytes)), int32(len(dbVerifierBytes))) != 1 ||
		subtle.ConstantTimeCompare(verifierBytes[:], dbVerifierBytes) != 1 {
		logger.Info("stored confirm verifier does not match provided one")
		return c.invalidToken(w, r)
	}

	user.PutConfirmSelector("")
	user.PutConfirmVerifier("")
	user.PutConfirmed(true)

	logger.Infof("user %s confirmed their account", user.GetPID())
	if err = c.Authboss.Config.Storage.Server.Save(r.Context(), user); err != nil {
		return err
	}

	ro := authboss.RedirectOptions{
		Code:         http.StatusTemporaryRedirect,
		Success:      "You have successfully confirmed your account.",
		RedirectPath: c.Authboss.Config.Paths.ConfirmOK,
	}
	return c.Authboss.Config.Core.Redirector.Redirect(w, r, ro)
}

func (c *Confirm) mailURL(token string) string {
	query := url.Values{FormValueConfirm: []string{token}}

	if len(c.Config.Mail.RootURL) != 0 {
		return fmt.Sprintf("%s?%s", c.Config.Mail.RootURL+"/confirm", query.Encode())
	}

	p := path.Join(c.Config.Paths.Mount, "confirm")
	return fmt.Sprintf("%s%s?%s", c.Config.Paths.RootURL, p, query.Encode())
}

func (c *Confirm) invalidToken(w http.ResponseWriter, r *http.Request) error {
	ro := authboss.RedirectOptions{
		Code:         http.StatusTemporaryRedirect,
		Failure:      "confirm token is invalid",
		RedirectPath: c.Authboss.Config.Paths.ConfirmNotOK,
	}
	return c.Authboss.Config.Core.Redirector.Redirect(w, r, ro)
}

// Middleware ensures that a user is confirmed, or else it will intercept the
// request and send them to the confirm page, this will load the user if he's
// not been loaded yet from the session.
//
// Panics if the user was not able to be loaded in order to allow a panic
// handler to show a nice error page, also panics if it failed to redirect
// for whatever reason.
func Middleware(ab *authboss.Authboss) func(http.Handler) http.Handler {
	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			user := ab.LoadCurrentUserP(&r)

			cu := authboss.MustBeConfirmable(user)
			if cu.GetConfirmed() {
				next.ServeHTTP(w, r)
				return
			}

			logger := ab.RequestLogger(r)
			logger.Infof("user %s prevented from accessing %s: not confirmed", user.GetPID(), r.URL.Path)
			ro := authboss.RedirectOptions{
				Code:         http.StatusTemporaryRedirect,
				Failure:      "Your account has not been confirmed, please check your e-mail.",
				RedirectPath: ab.Config.Paths.ConfirmNotOK,
			}
			if err := ab.Config.Core.Redirector.Redirect(w, r, ro); err != nil {
				logger.Errorf("error redirecting in confirm.Middleware: #%v", err)
			}
		})
	}
}

// GenerateConfirmCreds generates pieces needed for user confirm
// selector: hash of the first half of a 64 byte value
// (to be stored in the database and used in SELECT query)
// verifier: hash of the second half of a 64 byte value
// (to be stored in database but never used in SELECT query)
// token: the user-facing base64 encoded selector+verifier
func GenerateConfirmCreds() (selector, verifier, token string, err error) {
	rawToken := make([]byte, confirmTokenSize)
	if _, err = io.ReadFull(rand.Reader, rawToken); err != nil {
		return "", "", "", err
	}
	selectorBytes := sha512.Sum512(rawToken[:confirmTokenSplit])
	verifierBytes := sha512.Sum512(rawToken[confirmTokenSplit:])

	return base64.StdEncoding.EncodeToString(selectorBytes[:]),
		base64.StdEncoding.EncodeToString(verifierBytes[:]),
		base64.URLEncoding.EncodeToString(rawToken),
		nil
}
<|endoftext|>


<commit_before>package conf

import (
	"errors"
	"log"
	"net/url"
	"os"
	"path/filepath"
	"strconv"
)

// Config contains the application configuration settings
type Config struct {
	// Port gets the port number under which the site is being hosted.
	Port int

	// RequireSSL specifies whether to redirect all HTTP traffic to HTTPS. Default = true
	RequireSSL bool

	// UploadDriver is used to select which backend data store is used for file uploads.
	// Supported drivers: fs, s3
	UploadDriver string

	// UploadPath gets the path (full or relative) under which to store uploads.
	// When using Amazon S3, this should be set to the bucket name.
	UploadPath string

	// IsDevelopment defines whether to run the application in "development mode".
	// Development mode turns on additional features, such as logging, that may
	// not be desirable in a production environment.
	IsDevelopment bool

	// SecretKey is used for session authentication. Recommended to be 32 or 64 ASCII characters.
	SecretKey string

	// ApplicationTitle is used where the application name (title) is displayed on screen.
	ApplicationTitle string

	// HomeTitle is an optional heading displayed at the top of the gome screen.
	HomeTitle string

	// HomeImage is an optional heading image displayed beneath the HomeTitle.
	HomeImage string

	// DatabaseDriver gets which database/sql driver to use.
	// Supported drivers: postgres
	DatabaseDriver string

	// DatabaseUrl gets the url (or path, connection string, etc) to use with the associated
	// database driver when opening the database connection.
	DatabaseURL string
}

// Load reads the configuration file from the specified path
func Load(path string) *Config {
	c := Config{
		Port:             4000,
		RequireSSL:       true,
		UploadDriver:     "fs",
		UploadPath:       filepath.Join("data", "uploads"),
		IsDevelopment:    false,
		SecretKey:        "Secret123",
		ApplicationTitle: "GOMP: Go Meal Planner",
		HomeTitle:        "",
		HomeImage:        "",
		DatabaseDriver:   "postgres",
		DatabaseURL:      "",
	}

	// If environment variables are set, use them.
	loadEnv("PORT", &c.Port)
	loadEnv("GOMP_REQUIRE_SSL", &c.RequireSSL)
	loadEnv("GOMP_UPLOAD_DRIVER", &c.UploadDriver)
	loadEnv("GOMP_UPLOAD_PATH", &c.UploadPath)
	loadEnv("GOMP_IS_DEVELOPMENT", &c.IsDevelopment)
	loadEnv("GOMP_SECRET_KEY", &c.SecretKey)
	loadEnv("GOMP_APPLICATION_TITLE", &c.ApplicationTitle)
	loadEnv("GOMP_HOME_TITLE", &c.HomeTitle)
	loadEnv("GOMP_HOME_IMAGE", &c.HomeImage)
	loadEnv("DATABASE_DRIVER", &c.DatabaseDriver)
	loadEnv("DATABASE_URL", &c.DatabaseURL)

	if c.IsDevelopment {
		log.Printf("[config] Port=%d", c.Port)
		log.Printf("[config] RequireSSL=%d", c.RequireSSL)
		log.Printf("[config] UploadDriver=%s", c.UploadDriver)
		log.Printf("[config] UploadPath=%s", c.UploadPath)
		log.Printf("[config] IsDevelopment=%t", c.IsDevelopment)
		log.Printf("[config] SecretKey=%s", c.SecretKey)
		log.Printf("[config] ApplicationTitle=%s", c.ApplicationTitle)
		log.Printf("[config] HomeTitle=%s", c.HomeTitle)
		log.Printf("[config] HomeImage=%s", c.HomeImage)
		log.Printf("[config] DatabaseDriver=%s", c.DatabaseDriver)
		log.Printf("[config] DatabaseURL=%s", c.DatabaseURL)
	}

	return &c
}

// Validate checks whether the current configuration settings are valid.
func (c *Config) Validate() error {
	if c.Port <= 0 {
		return errors.New("PORT must be a positive integer")
	}

	if c.UploadDriver != "fs" && c.UploadDriver != "s3" {
		return errors.New("UPLOAD_DRIVER must be one of ('fs', 's3')")
	}

	if c.UploadPath == "" {
		return errors.New("UPLOAD_PATH must be specified")
	}

	if c.SecretKey == "" {
		return errors.New("GOMP_SECRET_KEY must be specified")
	}

	if c.ApplicationTitle == "" {
		return errors.New("GOMP_APPLICATION_TITLE must be specified")
	}

	if c.DatabaseDriver != "postgres" {
		return errors.New("DATABASE_DRIVER must be one of ('postgres')")
	}

	if c.DatabaseURL == "" {
		return errors.New("DATABASE_URL must be specified")
	}

	if _, err := url.Parse(c.DatabaseURL); err != nil {
		return errors.New("DATABASE_URL is invalid")
	}

	return nil
}

func loadEnv(name string, dest interface{}) {
	var err error
	if envStr := os.Getenv(name); envStr != "" {
		switch dest := dest.(type) {
		case *string:
			*dest = envStr
		case *int:
			if *dest, err = strconv.Atoi(envStr); err != nil {
				log.Fatalf("[config] Failed to convert %s environment variable to an integer. Value = %s, Error = %s",
					name, envStr, err)
			}
		case *bool:
			*dest = envStr != "0"
		}
	}
}
<commit_msg>Fixed a format string placeholder<commit_after>package conf

import (
	"errors"
	"log"
	"net/url"
	"os"
	"path/filepath"
	"strconv"
)

// Config contains the application configuration settings
type Config struct {
	// Port gets the port number under which the site is being hosted.
	Port int

	// RequireSSL specifies whether to redirect all HTTP traffic to HTTPS. Default = true
	RequireSSL bool

	// UploadDriver is used to select which backend data store is used for file uploads.
	// Supported drivers: fs, s3
	UploadDriver string

	// UploadPath gets the path (full or relative) under which to store uploads.
	// When using Amazon S3, this should be set to the bucket name.
	UploadPath string

	// IsDevelopment defines whether to run the application in "development mode".
	// Development mode turns on additional features, such as logging, that may
	// not be desirable in a production environment.
	IsDevelopment bool

	// SecretKey is used for session authentication. Recommended to be 32 or 64 ASCII characters.
	SecretKey string

	// ApplicationTitle is used where the application name (title) is displayed on screen.
	ApplicationTitle string

	// HomeTitle is an optional heading displayed at the top of the gome screen.
	HomeTitle string

	// HomeImage is an optional heading image displayed beneath the HomeTitle.
	HomeImage string

	// DatabaseDriver gets which database/sql driver to use.
	// Supported drivers: postgres
	DatabaseDriver string

	// DatabaseUrl gets the url (or path, connection string, etc) to use with the associated
	// database driver when opening the database connection.
	DatabaseURL string
}

// Load reads the configuration file from the specified path
func Load(path string) *Config {
	c := Config{
		Port:             4000,
		RequireSSL:       true,
		UploadDriver:     "fs",
		UploadPath:       filepath.Join("data", "uploads"),
		IsDevelopment:    false,
		SecretKey:        "Secret123",
		ApplicationTitle: "GOMP: Go Meal Planner",
		HomeTitle:        "",
		HomeImage:        "",
		DatabaseDriver:   "postgres",
		DatabaseURL:      "",
	}

	// If environment variables are set, use them.
	loadEnv("PORT", &c.Port)
	loadEnv("GOMP_REQUIRE_SSL", &c.RequireSSL)
	loadEnv("GOMP_UPLOAD_DRIVER", &c.UploadDriver)
	loadEnv("GOMP_UPLOAD_PATH", &c.UploadPath)
	loadEnv("GOMP_IS_DEVELOPMENT", &c.IsDevelopment)
	loadEnv("GOMP_SECRET_KEY", &c.SecretKey)
	loadEnv("GOMP_APPLICATION_TITLE", &c.ApplicationTitle)
	loadEnv("GOMP_HOME_TITLE", &c.HomeTitle)
	loadEnv("GOMP_HOME_IMAGE", &c.HomeImage)
	loadEnv("DATABASE_DRIVER", &c.DatabaseDriver)
	loadEnv("DATABASE_URL", &c.DatabaseURL)

	if c.IsDevelopment {
		log.Printf("[config] Port=%d", c.Port)
		log.Printf("[config] RequireSSL=%t", c.RequireSSL)
		log.Printf("[config] UploadDriver=%s", c.UploadDriver)
		log.Printf("[config] UploadPath=%s", c.UploadPath)
		log.Printf("[config] IsDevelopment=%t", c.IsDevelopment)
		log.Printf("[config] SecretKey=%s", c.SecretKey)
		log.Printf("[config] ApplicationTitle=%s", c.ApplicationTitle)
		log.Printf("[config] HomeTitle=%s", c.HomeTitle)
		log.Printf("[config] HomeImage=%s", c.HomeImage)
		log.Printf("[config] DatabaseDriver=%s", c.DatabaseDriver)
		log.Printf("[config] DatabaseURL=%s", c.DatabaseURL)
	}

	return &c
}

// Validate checks whether the current configuration settings are valid.
func (c *Config) Validate() error {
	if c.Port <= 0 {
		return errors.New("PORT must be a positive integer")
	}

	if c.UploadDriver != "fs" && c.UploadDriver != "s3" {
		return errors.New("UPLOAD_DRIVER must be one of ('fs', 's3')")
	}

	if c.UploadPath == "" {
		return errors.New("UPLOAD_PATH must be specified")
	}

	if c.SecretKey == "" {
		return errors.New("GOMP_SECRET_KEY must be specified")
	}

	if c.ApplicationTitle == "" {
		return errors.New("GOMP_APPLICATION_TITLE must be specified")
	}

	if c.DatabaseDriver != "postgres" {
		return errors.New("DATABASE_DRIVER must be one of ('postgres')")
	}

	if c.DatabaseURL == "" {
		return errors.New("DATABASE_URL must be specified")
	}

	if _, err := url.Parse(c.DatabaseURL); err != nil {
		return errors.New("DATABASE_URL is invalid")
	}

	return nil
}

func loadEnv(name string, dest interface{}) {
	var err error
	if envStr := os.Getenv(name); envStr != "" {
		switch dest := dest.(type) {
		case *string:
			*dest = envStr
		case *int:
			if *dest, err = strconv.Atoi(envStr); err != nil {
				log.Fatalf("[config] Failed to convert %s environment variable to an integer. Value = %s, Error = %s",
					name, envStr, err)
			}
		case *bool:
			*dest = envStr != "0"
		}
	}
}
<|endoftext|>


<commit_before>// Copyright 2015 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 influxdb

import (
	"fmt"
	"net/url"
	"strings"
	"sync"
	"time"

	influxdb_common "k8s.io/heapster/common/influxdb"
	"k8s.io/heapster/metrics/core"

	"github.com/golang/glog"
	influxdb "github.com/influxdata/influxdb/client"
)

type influxdbSink struct {
	client influxdb_common.InfluxdbClient
	sync.RWMutex
	c        influxdb_common.InfluxdbConfig
	dbExists bool
}

var influxdbBlacklistLabels = map[string]struct{}{
	core.LabelPodNamespaceUID.Key: {},
	core.LabelPodId.Key:           {},
	core.LabelHostname.Key:        {},
	core.LabelHostID.Key:          {},
}

const (
	// Value Field name
	valueField = "value"
	// Event special tags
	dbNotFoundError = "database not found"

	// Maximum number of influxdb Points to be sent in one batch.
	maxSendBatchSize = 10000
)

func (sink *influxdbSink) resetConnection() {
	glog.Infof("Influxdb connection reset")
	sink.dbExists = false
	sink.client = nil
}

func (sink *influxdbSink) ExportData(dataBatch *core.DataBatch) {
	sink.Lock()
	defer sink.Unlock()

	dataPoints := make([]influxdb.Point, 0, 0)
	for _, metricSet := range dataBatch.MetricSets {
		for metricName, metricValue := range metricSet.MetricValues {

			var value interface{}
			if core.ValueInt64 == metricValue.ValueType {
				value = metricValue.IntValue
			} else if core.ValueFloat == metricValue.ValueType {
				value = float64(metricValue.FloatValue)
			} else {
				continue
			}

			// Prepare measurement without fields
			fieldName := "value"
			measurementName := metricName
			if sink.c.WithFields {
				// Prepare measurement and field names
				serieName := strings.SplitN(metricName, "/", 2)
				measurementName = serieName[0]
				if len(serieName) > 1 {
					fieldName = serieName[1]
				}
			}

			point := influxdb.Point{
				Measurement: measurementName,
				Tags:        make(map[string]string, len(metricSet.Labels)),
				Fields: map[string]interface{}{
					fieldName: value,
				},
				Time: dataBatch.Timestamp.UTC(),
			}
			for key, value := range metricSet.Labels {
				if _, exists := influxdbBlacklistLabels[key]; !exists {
					if value != "" {
						point.Tags[key] = value
					}
				}
			}

			point.Tags["cluster_name"] = sink.c.ClusterName

			dataPoints = append(dataPoints, point)
			if len(dataPoints) >= maxSendBatchSize {
				sink.sendData(dataPoints)
				dataPoints = make([]influxdb.Point, 0, 0)
			}
		}

		for _, labeledMetric := range metricSet.LabeledMetrics {

			var value interface{}
			if core.ValueInt64 == labeledMetric.ValueType {
				value = labeledMetric.IntValue
			} else if core.ValueFloat == labeledMetric.ValueType {
				value = float64(labeledMetric.FloatValue)
			} else {
				continue
			}

			// Prepare measurement without fields
			fieldName := "value"
			measurementName := labeledMetric.Name
			if sink.c.WithFields {
				// Prepare measurement and field names
				serieName := strings.SplitN(labeledMetric.Name, "/", 2)
				measurementName = serieName[0]
				if len(serieName) > 1 {
					fieldName = serieName[1]
				}
			}

			point := influxdb.Point{
				Measurement: measurementName,
				Tags:        make(map[string]string, len(metricSet.Labels)+len(labeledMetric.Labels)),
				Fields: map[string]interface{}{
					fieldName: value,
				},
				Time: dataBatch.Timestamp.UTC(),
			}

			for key, value := range metricSet.Labels {
				if value != "" {
					point.Tags[key] = value
				}
			}
			for key, value := range labeledMetric.Labels {
				if _, exists := influxdbBlacklistLabels[key]; !exists {
					if value != "" {
						point.Tags[key] = value
					}
				}
			}
			point.Tags["cluster_name"] = sink.c.ClusterName

			dataPoints = append(dataPoints, point)
			if len(dataPoints) >= maxSendBatchSize {
				sink.sendData(dataPoints)
				dataPoints = make([]influxdb.Point, 0, 0)
			}
		}
	}
	if len(dataPoints) >= 0 {
		sink.sendData(dataPoints)
	}
}

func (sink *influxdbSink) sendData(dataPoints []influxdb.Point) {
	if err := sink.createDatabase(); err != nil {
		glog.Errorf("Failed to create influxdb: %v", err)
		return
	}
	bp := influxdb.BatchPoints{
		Points:          dataPoints,
		Database:        sink.c.DbName,
		RetentionPolicy: "default",
	}

	start := time.Now()
	if _, err := sink.client.Write(bp); err != nil {
		glog.Errorf("InfluxDB write failed: %v", err)
		if strings.Contains(err.Error(), dbNotFoundError) {
			sink.resetConnection()
		} else if _, _, err := sink.client.Ping(); err != nil {
			glog.Errorf("InfluxDB ping failed: %v", err)
			sink.resetConnection()
		}
		return
	}
	end := time.Now()
	glog.V(4).Infof("Exported %d data to influxDB in %s", len(dataPoints), end.Sub(start))
}

func (sink *influxdbSink) Name() string {
	return "InfluxDB Sink"
}

func (sink *influxdbSink) Stop() {
	// nothing needs to be done.
}

func (sink *influxdbSink) ensureClient() error {
	if sink.client == nil {
		client, err := influxdb_common.NewClient(sink.c)
		if err != nil {
			return err
		}
		sink.client = client
	}

	return nil
}

func (sink *influxdbSink) createDatabase() error {
	if err := sink.ensureClient(); err != nil {
		return err
	}

	if sink.dbExists {
		return nil
	}
	q := influxdb.Query{
		Command: fmt.Sprintf(`CREATE DATABASE %s WITH NAME "default"`, sink.c.DbName),
	}

	if resp, err := sink.client.Query(q); err != nil {
		if !(resp != nil && resp.Err != nil && strings.Contains(resp.Err.Error(), "already exists")) {
			err := sink.createRetentionPolicy()
			if err != nil {
				return err
			}
		}
	}

	sink.dbExists = true
	glog.Infof("Created database %q on influxDB server at %q", sink.c.DbName, sink.c.Host)
	return nil
}

func (sink *influxdbSink) createRetentionPolicy() error {
	q := influxdb.Query{
		Command: fmt.Sprintf(`CREATE RETENTION POLICY "default" ON %s DURATION %s REPLICATION 1 DEFAULT`, sink.c.DbName, sink.c.RetentionPolicy),
	}

	if resp, err := sink.client.Query(q); err != nil {
		if !(resp != nil && resp.Err != nil) {
			return fmt.Errorf("Retention Policy creation failed: %v", err)
		}
	}

	glog.Infof("Created retention policy 'default' in database %q on influxDB server at %q", sink.c.DbName, sink.c.Host)
	return nil
}

// Returns a thread-compatible implementation of influxdb interactions.
func new(c influxdb_common.InfluxdbConfig) core.DataSink {
	client, err := influxdb_common.NewClient(c)
	if err != nil {
		glog.Errorf("issues while creating an InfluxDB sink: %v, will retry on use", err)
	}
	return &influxdbSink{
		client: client, // can be nil
		c:      c,
	}
}

func CreateInfluxdbSink(uri *url.URL) (core.DataSink, error) {
	config, err := influxdb_common.BuildConfig(uri)
	if err != nil {
		return nil, err
	}
	sink := new(*config)
	glog.Infof("created influxdb sink with options: host:%s user:%s db:%s", config.Host, config.User, config.DbName)
	return sink, nil
}
<commit_msg>fix remove blacklist labels<commit_after>// Copyright 2015 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 influxdb

import (
	"fmt"
	"net/url"
	"strings"
	"sync"
	"time"

	influxdb_common "k8s.io/heapster/common/influxdb"
	"k8s.io/heapster/metrics/core"

	"github.com/golang/glog"
	influxdb "github.com/influxdata/influxdb/client"
)

type influxdbSink struct {
	client influxdb_common.InfluxdbClient
	sync.RWMutex
	c        influxdb_common.InfluxdbConfig
	dbExists bool
}

var influxdbBlacklistLabels = map[string]struct{}{
	core.LabelPodNamespaceUID.Key: {},
	core.LabelPodId.Key:           {},
	core.LabelHostname.Key:        {},
	core.LabelHostID.Key:          {},
}

const (
	// Value Field name
	valueField = "value"
	// Event special tags
	dbNotFoundError = "database not found"

	// Maximum number of influxdb Points to be sent in one batch.
	maxSendBatchSize = 10000
)

func (sink *influxdbSink) resetConnection() {
	glog.Infof("Influxdb connection reset")
	sink.dbExists = false
	sink.client = nil
}

func (sink *influxdbSink) ExportData(dataBatch *core.DataBatch) {
	sink.Lock()
	defer sink.Unlock()

	dataPoints := make([]influxdb.Point, 0, 0)
	for _, metricSet := range dataBatch.MetricSets {
		for metricName, metricValue := range metricSet.MetricValues {

			var value interface{}
			if core.ValueInt64 == metricValue.ValueType {
				value = metricValue.IntValue
			} else if core.ValueFloat == metricValue.ValueType {
				value = float64(metricValue.FloatValue)
			} else {
				continue
			}

			// Prepare measurement without fields
			fieldName := "value"
			measurementName := metricName
			if sink.c.WithFields {
				// Prepare measurement and field names
				serieName := strings.SplitN(metricName, "/", 2)
				measurementName = serieName[0]
				if len(serieName) > 1 {
					fieldName = serieName[1]
				}
			}

			point := influxdb.Point{
				Measurement: measurementName,
				Tags:        make(map[string]string, len(metricSet.Labels)),
				Fields: map[string]interface{}{
					fieldName: value,
				},
				Time: dataBatch.Timestamp.UTC(),
			}
			for key, value := range metricSet.Labels {
				if _, exists := influxdbBlacklistLabels[key]; !exists {
					if value != "" {
						point.Tags[key] = value
					}
				}
			}

			point.Tags["cluster_name"] = sink.c.ClusterName

			dataPoints = append(dataPoints, point)
			if len(dataPoints) >= maxSendBatchSize {
				sink.sendData(dataPoints)
				dataPoints = make([]influxdb.Point, 0, 0)
			}
		}

		for _, labeledMetric := range metricSet.LabeledMetrics {

			var value interface{}
			if core.ValueInt64 == labeledMetric.ValueType {
				value = labeledMetric.IntValue
			} else if core.ValueFloat == labeledMetric.ValueType {
				value = float64(labeledMetric.FloatValue)
			} else {
				continue
			}

			// Prepare measurement without fields
			fieldName := "value"
			measurementName := labeledMetric.Name
			if sink.c.WithFields {
				// Prepare measurement and field names
				serieName := strings.SplitN(labeledMetric.Name, "/", 2)
				measurementName = serieName[0]
				if len(serieName) > 1 {
					fieldName = serieName[1]
				}
			}

			point := influxdb.Point{
				Measurement: measurementName,
				Tags:        make(map[string]string, len(metricSet.Labels)+len(labeledMetric.Labels)),
				Fields: map[string]interface{}{
					fieldName: value,
				},
				Time: dataBatch.Timestamp.UTC(),
			}

			for key, value := range metricSet.Labels {
				if _, exists := influxdbBlacklistLabels[key]; !exists {
					if value != "" {
						point.Tags[key] = value
					}
				}
			}
			for key, value := range labeledMetric.Labels {
				if _, exists := influxdbBlacklistLabels[key]; !exists {
					if value != "" {
						point.Tags[key] = value
					}
				}
			}
			point.Tags["cluster_name"] = sink.c.ClusterName

			dataPoints = append(dataPoints, point)
			if len(dataPoints) >= maxSendBatchSize {
				sink.sendData(dataPoints)
				dataPoints = make([]influxdb.Point, 0, 0)
			}
		}
	}
	if len(dataPoints) >= 0 {
		sink.sendData(dataPoints)
	}
}

func (sink *influxdbSink) sendData(dataPoints []influxdb.Point) {
	if err := sink.createDatabase(); err != nil {
		glog.Errorf("Failed to create influxdb: %v", err)
		return
	}
	bp := influxdb.BatchPoints{
		Points:          dataPoints,
		Database:        sink.c.DbName,
		RetentionPolicy: "default",
	}

	start := time.Now()
	if _, err := sink.client.Write(bp); err != nil {
		glog.Errorf("InfluxDB write failed: %v", err)
		if strings.Contains(err.Error(), dbNotFoundError) {
			sink.resetConnection()
		} else if _, _, err := sink.client.Ping(); err != nil {
			glog.Errorf("InfluxDB ping failed: %v", err)
			sink.resetConnection()
		}
		return
	}
	end := time.Now()
	glog.V(4).Infof("Exported %d data to influxDB in %s", len(dataPoints), end.Sub(start))
}

func (sink *influxdbSink) Name() string {
	return "InfluxDB Sink"
}

func (sink *influxdbSink) Stop() {
	// nothing needs to be done.
}

func (sink *influxdbSink) ensureClient() error {
	if sink.client == nil {
		client, err := influxdb_common.NewClient(sink.c)
		if err != nil {
			return err
		}
		sink.client = client
	}

	return nil
}

func (sink *influxdbSink) createDatabase() error {
	if err := sink.ensureClient(); err != nil {
		return err
	}

	if sink.dbExists {
		return nil
	}
	q := influxdb.Query{
		Command: fmt.Sprintf(`CREATE DATABASE %s WITH NAME "default"`, sink.c.DbName),
	}

	if resp, err := sink.client.Query(q); err != nil {
		if !(resp != nil && resp.Err != nil && strings.Contains(resp.Err.Error(), "already exists")) {
			err := sink.createRetentionPolicy()
			if err != nil {
				return err
			}
		}
	}

	sink.dbExists = true
	glog.Infof("Created database %q on influxDB server at %q", sink.c.DbName, sink.c.Host)
	return nil
}

func (sink *influxdbSink) createRetentionPolicy() error {
	q := influxdb.Query{
		Command: fmt.Sprintf(`CREATE RETENTION POLICY "default" ON %s DURATION %s REPLICATION 1 DEFAULT`, sink.c.DbName, sink.c.RetentionPolicy),
	}

	if resp, err := sink.client.Query(q); err != nil {
		if !(resp != nil && resp.Err != nil) {
			return fmt.Errorf("Retention Policy creation failed: %v", err)
		}
	}

	glog.Infof("Created retention policy 'default' in database %q on influxDB server at %q", sink.c.DbName, sink.c.Host)
	return nil
}

// Returns a thread-compatible implementation of influxdb interactions.
func new(c influxdb_common.InfluxdbConfig) core.DataSink {
	client, err := influxdb_common.NewClient(c)
	if err != nil {
		glog.Errorf("issues while creating an InfluxDB sink: %v, will retry on use", err)
	}
	return &influxdbSink{
		client: client, // can be nil
		c:      c,
	}
}

func CreateInfluxdbSink(uri *url.URL) (core.DataSink, error) {
	config, err := influxdb_common.BuildConfig(uri)
	if err != nil {
		return nil, err
	}
	sink := new(*config)
	glog.Infof("created influxdb sink with options: host:%s user:%s db:%s", config.Host, config.User, config.DbName)
	return sink, nil
}
<|endoftext|>


<commit_before>package rate

import (
	"context"
	"errors"
	"math"
	"time"

	"github.com/lukechampine/stm"
	"github.com/lukechampine/stm/stmutil"
)

type numTokens = int

type Limiter struct {
	max     *stm.Var
	cur     *stm.Var
	lastAdd *stm.Var
	rate    Limit
}

const Inf = Limit(math.MaxFloat64)

type Limit float64

func (l Limit) interval() time.Duration {
	return time.Duration(Limit(1*time.Second) / l)
}

func Every(interval time.Duration) Limit {
	if interval == 0 {
		return Inf
	}
	return Limit(time.Second / interval)
}

func NewLimiter(rate Limit, burst numTokens) *Limiter {
	rl := &Limiter{
		max:     stm.NewVar(burst),
		cur:     stm.NewVar(burst),
		lastAdd: stm.NewVar(time.Now()),
		rate:    rate,
	}
	if rate != Inf {
		go rl.tokenGenerator(rate.interval())
	}
	return rl
}

func (rl *Limiter) tokenGenerator(interval time.Duration) {
	for {
		lastAdd := stm.AtomicGet(rl.lastAdd).(time.Time)
		time.Sleep(time.Until(lastAdd.Add(interval)))
		now := time.Now()
		available := numTokens(now.Sub(lastAdd) / interval)
		if available < 1 {
			continue
		}
		stm.Atomically(func(tx *stm.Tx) {
			cur := tx.Get(rl.cur).(numTokens)
			max := tx.Get(rl.max).(numTokens)
			tx.Assert(cur < max)
			newCur := cur + available
			if newCur > max {
				newCur = max
			}
			if newCur != cur {
				tx.Set(rl.cur, newCur)
			}
			tx.Set(rl.lastAdd, now)
		})
	}
}

func (rl *Limiter) Allow() bool {
	return rl.AllowN(1)
}

func (rl *Limiter) AllowN(n numTokens) bool {
	return stm.Atomically(func(tx *stm.Tx) {
		tx.Return(rl.takeTokens(tx, n))
	}).(bool)
}

func (rl *Limiter) AllowStm(tx *stm.Tx) bool {
	return rl.takeTokens(tx, 1)
}

func (rl *Limiter) takeTokens(tx *stm.Tx, n numTokens) bool {
	if rl.rate == Inf {
		return true
	}
	cur := tx.Get(rl.cur).(numTokens)
	if cur >= n {
		tx.Set(rl.cur, cur-n)
		return true
	}
	return false
}

func (rl *Limiter) Wait(ctx context.Context) error {
	return rl.WaitN(ctx, 1)
}

func (rl *Limiter) WaitN(ctx context.Context, n int) error {
	ctxDone, cancel := stmutil.ContextDoneVar(ctx)
	defer cancel()
	if err := stm.Atomically(func(tx *stm.Tx) {
		if tx.Get(ctxDone).(bool) {
			tx.Return(ctx.Err())
		}
		if rl.takeTokens(tx, n) {
			tx.Return(nil)
		}
		if n > tx.Get(rl.max).(numTokens) {
			tx.Return(errors.New("burst exceeded"))
		}
		if dl, ok := ctx.Deadline(); ok {
			if tx.Get(rl.cur).(numTokens)+numTokens(dl.Sub(tx.Get(rl.lastAdd).(time.Time))/rl.rate.interval()) < n {
				tx.Return(context.DeadlineExceeded)
			}
		}
		tx.Retry()
	}); err != nil {
		return err.(error)
	}
	return nil

}
<commit_msg>Fix rate.TestLongRunningQPS<commit_after>package rate

import (
	"context"
	"errors"
	"math"
	"time"

	"github.com/lukechampine/stm"
	"github.com/lukechampine/stm/stmutil"
)

type numTokens = int

type Limiter struct {
	max     *stm.Var
	cur     *stm.Var
	lastAdd *stm.Var
	rate    Limit
}

const Inf = Limit(math.MaxFloat64)

type Limit float64

func (l Limit) interval() time.Duration {
	return time.Duration(Limit(1*time.Second) / l)
}

func Every(interval time.Duration) Limit {
	if interval == 0 {
		return Inf
	}
	return Limit(time.Second / interval)
}

func NewLimiter(rate Limit, burst numTokens) *Limiter {
	rl := &Limiter{
		max:     stm.NewVar(burst),
		cur:     stm.NewVar(burst),
		lastAdd: stm.NewVar(time.Now()),
		rate:    rate,
	}
	if rate != Inf {
		go rl.tokenGenerator(rate.interval())
	}
	return rl
}

func (rl *Limiter) tokenGenerator(interval time.Duration) {
	for {
		lastAdd := stm.AtomicGet(rl.lastAdd).(time.Time)
		time.Sleep(time.Until(lastAdd.Add(interval)))
		now := time.Now()
		available := numTokens(now.Sub(lastAdd) / interval)
		if available < 1 {
			continue
		}
		stm.Atomically(func(tx *stm.Tx) {
			cur := tx.Get(rl.cur).(numTokens)
			max := tx.Get(rl.max).(numTokens)
			tx.Assert(cur < max)
			newCur := cur + available
			if newCur > max {
				newCur = max
			}
			if newCur != cur {
				tx.Set(rl.cur, newCur)
			}
			tx.Set(rl.lastAdd, lastAdd.Add(interval*time.Duration(available)))
		})
	}
}

func (rl *Limiter) Allow() bool {
	return rl.AllowN(1)
}

func (rl *Limiter) AllowN(n numTokens) bool {
	return stm.Atomically(func(tx *stm.Tx) {
		tx.Return(rl.takeTokens(tx, n))
	}).(bool)
}

func (rl *Limiter) AllowStm(tx *stm.Tx) bool {
	return rl.takeTokens(tx, 1)
}

func (rl *Limiter) takeTokens(tx *stm.Tx, n numTokens) bool {
	if rl.rate == Inf {
		return true
	}
	cur := tx.Get(rl.cur).(numTokens)
	if cur >= n {
		tx.Set(rl.cur, cur-n)
		return true
	}
	return false
}

func (rl *Limiter) Wait(ctx context.Context) error {
	return rl.WaitN(ctx, 1)
}

func (rl *Limiter) WaitN(ctx context.Context, n int) error {
	ctxDone, cancel := stmutil.ContextDoneVar(ctx)
	defer cancel()
	if err := stm.Atomically(func(tx *stm.Tx) {
		if tx.Get(ctxDone).(bool) {
			tx.Return(ctx.Err())
		}
		if rl.takeTokens(tx, n) {
			tx.Return(nil)
		}
		if n > tx.Get(rl.max).(numTokens) {
			tx.Return(errors.New("burst exceeded"))
		}
		if dl, ok := ctx.Deadline(); ok {
			if tx.Get(rl.cur).(numTokens)+numTokens(dl.Sub(tx.Get(rl.lastAdd).(time.Time))/rl.rate.interval()) < n {
				tx.Return(context.DeadlineExceeded)
			}
		}
		tx.Retry()
	}); err != nil {
		return err.(error)
	}
	return nil

}
<|endoftext|>


<commit_before>package zygo
func init() { GITLASTTAG = "v4.8.3"; GITLASTCOMMIT = "2badbba3e494eb9731ae80f8358090474d079064" }
<commit_msg>atg. latest<commit_after>package zygo
func init() { GITLASTTAG = "v4.8.3"; GITLASTCOMMIT = "1433a9c89896798a1f0dac11223fc5dc5dcad874" }
<|endoftext|>


<commit_before>package gohm

import (
	"bytes"
	"context"
	"net/http"
	"sync/atomic"
	"time"
)

// responseWriter must behave exactly like http.ResponseWriter, yet store up response until query
// complete and flush invoked.
type responseWriter struct {
	http.ResponseWriter
	header        http.Header
	body          bytes.Buffer
	size          int64
	status        int
	statusWritten bool
	errorMessage  string
	begin, end    time.Time
}

func (rw *responseWriter) Header() http.Header {
	m := rw.header
	if m == nil {
		m = make(http.Header)
		rw.header = m
	}
	return m
}

func (rw *responseWriter) Write(blob []byte) (int, error) {
	return rw.body.Write(blob)
}

func (rw *responseWriter) WriteHeader(status int) {
	rw.status = status
	rw.statusWritten = true
}

// update responseWriter then enqueue status and message to be send to client
func (rw *responseWriter) error(message string, status int) {
	rw.errorMessage = message
	rw.status = status
	Error(rw, rw.errorMessage, rw.status)
}

func (rw *responseWriter) flush() error {
	// write header
	header := rw.ResponseWriter.Header()
	for key, values := range rw.header {
		for _, value := range values {
			header.Add(key, value)
		}
	}

	// write status
	if !rw.statusWritten {
		rw.status = http.StatusOK
	}
	rw.ResponseWriter.WriteHeader(rw.status)

	// write response
	var err error

	// NOTE: Apache Common Log Format size excludes HTTP headers
	rw.size, err = rw.body.WriteTo(rw.ResponseWriter)
	return err
}

// New returns a new http.Handler that calls the specified next http.Handler, and performs the
// requested operations before and after the downstream handler as specified by the gohm.Config
// structure passed to it.
//
// It receives a gohm.Config struct rather than a pointer to one, so users less likely to consider
// modification after creating the http.Handler.
//
//	const staticTimeout = time.Second // Used to control how long it takes to serve a static file.
//
//	var (
//		// Will store statistics counters for status codes 1xx, 2xx, 3xx, 4xx, 5xx, as well as a
//		// counter for all responses
//		counters gohm.Counters
//
//		// Used to dynamically control log level of HTTP logging.  After handler created, this must
//		// be accessed using the sync/atomic package.
//		logBitmask = gohm.LogStatusErrors
//
//		// Determines HTTP log format
//		logFormat = "{http-CLIENT-IP} {client-ip} [{end}] \"{method} {uri} {proto}\" {status} {bytes} {duration} {message}"
//	)
//
//	func main() {
//
//		h := http.StripPrefix("/static/", http.FileServer(http.Dir("static")))
//
//		h = gohm.WithGzip(h) // gzip response if client accepts gzip encoding
//
//		// gohm was designed to wrap other http.Handler functions.
//		h = gohm.New(h, gohm.Config{
//			Counters:   &counters,   // pointer given so counters can be collected and optionally reset
//			LogBitmask: &logBitmask, // pointer given so bitmask can be updated using sync/atomic
//			LogFormat:  logFormat,
//			LogWriter:  os.Stderr,
//			Timeout:    staticTimeout,
//		})
//
//		http.Handle("/static/", h)
//		log.Fatal(http.ListenAndServe(":8080", nil))
//	}
func New(next http.Handler, config Config) http.Handler {
	var emitters []func(*responseWriter, *http.Request, *bytes.Buffer)

	if config.LogWriter != nil {
		if config.LogBitmask == nil {
			// Set a default bitmask to log all requests
			logBitmask := LogStatusAll
			config.LogBitmask = &logBitmask
		}
		if config.LogFormat == "" {
			// Set a default log line format
			config.LogFormat = DefaultLogFormat
		}
		emitters = compileFormat(config.LogFormat)
	}

	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// Create a responseWriter to pass to next.ServeHTTP and collect downstream
		// handler's response to query.  It will eventually be used to flush to the client,
		// assuming neither the handler panics, nor the client connection is detected to be
		// closed.
		rw := &responseWriter{ResponseWriter: w}

		var ctx context.Context

		// Create a couple of channels to detect one of 3 ways to exit this handler.
		serverCompleted := make(chan struct{})
		serverPanicked := make(chan string, 1)

		if config.Timeout > 0 {
			// Adding a timeout to a request context spins off a goroutine that will
			// invoke the specified cancel function for us after the timeout has
			// elapsed.  Invoking the cancel function causes the context's Done channel
			// to close.  Detecting timeout is done by waiting for context.Done() to close.
			ctx, _ = context.WithTimeout(r.Context(), config.Timeout)
		} else {
			// When no timeout given, we still need a mechanism to track context
			// cancellation so this handler can detect when client has closed its
			// connection.
			ctx, _ = context.WithCancel(r.Context())
		}
		r = r.WithContext(ctx)

		if config.LogWriter != nil {
			rw.begin = time.Now()
		}

		// We must invoke downstream handler in separate goroutine in order to ensure this
		// handler only responds to one of the three events below, whichever event takes
		// place first.
		go serveWithPanicProtection(rw, r, next, serverCompleted, serverPanicked)

		// Wait for the first of either of 3 events:
		//   * serveComplete: the next.ServeHTTP method completed normally (possibly even
		//     with an erroneous status code).
		//   * servePanicked: the next.ServeHTTP method failed to complete, and panicked
		//     instead with a text message.
		//   * context is done: triggered when timeout or client disconnect.
		select {

		case <-serverCompleted:
			// break

		case text := <-serverPanicked:
			if config.AllowPanics {
				panic(text) // do not need to tell downstream to cancel, because it already panicked.
			}
			rw.error(text, http.StatusInternalServerError)

		case <-ctx.Done():
			// we'll create a new rw that downstream handler doesn't have access to so it cannot
			// mutate it.
			rw = &responseWriter{ResponseWriter: w}

			// the context was canceled; where ctx.Err() will say why
			// 503 (this is what http.TimeoutHandler returns)
			rw.error(ctx.Err().Error(), http.StatusServiceUnavailable)

		}

		if err := rw.flush(); err != nil {
			// cannot write responseWriter's contents to http.ResponseWriter
			rw.errorMessage = err.Error()
			rw.status = http.StatusInternalServerError
			// no use emitting error message to client when cannot send original payload back
		}

		statusClass := rw.status / 100

		// Update status counters
		if config.Counters != nil {
			atomic.AddUint64(&config.Counters.counters[0], 1)           // all
			atomic.AddUint64(&config.Counters.counters[statusClass], 1) // 1xx, 2xx, 3xx, 4xx, 5xx
		}

		// Update log
		if config.LogWriter != nil {
			var bit uint32 = 1 << uint32(statusClass-1)

			if (atomic.LoadUint32(config.LogBitmask))&bit > 0 {
				rw.end = time.Now()

				buf := bytes.NewBuffer(make([]byte, 0, 128))
				for _, emitter := range emitters {
					emitter(rw, r, buf)
				}
				_, _ = buf.WriteTo(config.LogWriter)
			}
		}
	})
}
<commit_msg>new response writer has begin time set<commit_after>package gohm

import (
	"bytes"
	"context"
	"net/http"
	"sync/atomic"
	"time"
)

// responseWriter must behave exactly like http.ResponseWriter, yet store up response until query
// complete and flush invoked.
type responseWriter struct {
	http.ResponseWriter
	header        http.Header
	body          bytes.Buffer
	size          int64
	status        int
	statusWritten bool
	errorMessage  string
	begin, end    time.Time
}

func (rw *responseWriter) Header() http.Header {
	m := rw.header
	if m == nil {
		m = make(http.Header)
		rw.header = m
	}
	return m
}

func (rw *responseWriter) Write(blob []byte) (int, error) {
	return rw.body.Write(blob)
}

func (rw *responseWriter) WriteHeader(status int) {
	rw.status = status
	rw.statusWritten = true
}

// update responseWriter then enqueue status and message to be send to client
func (rw *responseWriter) error(message string, status int) {
	rw.errorMessage = message
	rw.status = status
	Error(rw, rw.errorMessage, rw.status)
}

func (rw *responseWriter) flush() error {
	// write header
	header := rw.ResponseWriter.Header()
	for key, values := range rw.header {
		for _, value := range values {
			header.Add(key, value)
		}
	}

	// write status
	if !rw.statusWritten {
		rw.status = http.StatusOK
	}
	rw.ResponseWriter.WriteHeader(rw.status)

	// write response
	var err error

	// NOTE: Apache Common Log Format size excludes HTTP headers
	rw.size, err = rw.body.WriteTo(rw.ResponseWriter)
	return err
}

// New returns a new http.Handler that calls the specified next http.Handler, and performs the
// requested operations before and after the downstream handler as specified by the gohm.Config
// structure passed to it.
//
// It receives a gohm.Config struct rather than a pointer to one, so users less likely to consider
// modification after creating the http.Handler.
//
//	const staticTimeout = time.Second // Used to control how long it takes to serve a static file.
//
//	var (
//		// Will store statistics counters for status codes 1xx, 2xx, 3xx, 4xx, 5xx, as well as a
//		// counter for all responses
//		counters gohm.Counters
//
//		// Used to dynamically control log level of HTTP logging.  After handler created, this must
//		// be accessed using the sync/atomic package.
//		logBitmask = gohm.LogStatusErrors
//
//		// Determines HTTP log format
//		logFormat = "{http-CLIENT-IP} {client-ip} [{end}] \"{method} {uri} {proto}\" {status} {bytes} {duration} {message}"
//	)
//
//	func main() {
//
//		h := http.StripPrefix("/static/", http.FileServer(http.Dir("static")))
//
//		h = gohm.WithGzip(h) // gzip response if client accepts gzip encoding
//
//		// gohm was designed to wrap other http.Handler functions.
//		h = gohm.New(h, gohm.Config{
//			Counters:   &counters,   // pointer given so counters can be collected and optionally reset
//			LogBitmask: &logBitmask, // pointer given so bitmask can be updated using sync/atomic
//			LogFormat:  logFormat,
//			LogWriter:  os.Stderr,
//			Timeout:    staticTimeout,
//		})
//
//		http.Handle("/static/", h)
//		log.Fatal(http.ListenAndServe(":8080", nil))
//	}
func New(next http.Handler, config Config) http.Handler {
	var emitters []func(*responseWriter, *http.Request, *bytes.Buffer)

	if config.LogWriter != nil {
		if config.LogBitmask == nil {
			// Set a default bitmask to log all requests
			logBitmask := LogStatusAll
			config.LogBitmask = &logBitmask
		}
		if config.LogFormat == "" {
			// Set a default log line format
			config.LogFormat = DefaultLogFormat
		}
		emitters = compileFormat(config.LogFormat)
	}

	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// Create a responseWriter to pass to next.ServeHTTP and collect downstream
		// handler's response to query.  It will eventually be used to flush to the client,
		// assuming neither the handler panics, nor the client connection is detected to be
		// closed.
		rw := &responseWriter{ResponseWriter: w}

		var ctx context.Context

		// Create a couple of channels to detect one of 3 ways to exit this handler.
		serverCompleted := make(chan struct{})
		serverPanicked := make(chan string, 1)

		if config.Timeout > 0 {
			// Adding a timeout to a request context spins off a goroutine that will
			// invoke the specified cancel function for us after the timeout has
			// elapsed.  Invoking the cancel function causes the context's Done channel
			// to close.  Detecting timeout is done by waiting for context.Done() to close.
			ctx, _ = context.WithTimeout(r.Context(), config.Timeout)
		} else {
			// When no timeout given, we still need a mechanism to track context
			// cancellation so this handler can detect when client has closed its
			// connection.
			ctx, _ = context.WithCancel(r.Context())
		}
		r = r.WithContext(ctx)

		if config.LogWriter != nil {
			rw.begin = time.Now()
		}

		// We must invoke downstream handler in separate goroutine in order to ensure this
		// handler only responds to one of the three events below, whichever event takes
		// place first.
		go serveWithPanicProtection(rw, r, next, serverCompleted, serverPanicked)

		// Wait for the first of either of 3 events:
		//   * serveComplete: the next.ServeHTTP method completed normally (possibly even
		//     with an erroneous status code).
		//   * servePanicked: the next.ServeHTTP method failed to complete, and panicked
		//     instead with a text message.
		//   * context is done: triggered when timeout or client disconnect.
		select {

		case <-serverCompleted:
			// break

		case text := <-serverPanicked:
			if config.AllowPanics {
				panic(text) // do not need to tell downstream to cancel, because it already panicked.
			}
			rw.error(text, http.StatusInternalServerError)

		case <-ctx.Done():
			// we'll create a new rw that downstream handler doesn't have access to so it cannot
			// mutate it.
			rw = &responseWriter{ResponseWriter: w, begin: rw.begin}

			// the context was canceled; where ctx.Err() will say why
			// 503 (this is what http.TimeoutHandler returns)
			rw.error(ctx.Err().Error(), http.StatusServiceUnavailable)

		}

		if err := rw.flush(); err != nil {
			// cannot write responseWriter's contents to http.ResponseWriter
			rw.errorMessage = err.Error()
			rw.status = http.StatusInternalServerError
			// no use emitting error message to client when cannot send original payload back
		}

		statusClass := rw.status / 100

		// Update status counters
		if config.Counters != nil {
			atomic.AddUint64(&config.Counters.counters[0], 1)           // all
			atomic.AddUint64(&config.Counters.counters[statusClass], 1) // 1xx, 2xx, 3xx, 4xx, 5xx
		}

		// Update log
		if config.LogWriter != nil {
			var bit uint32 = 1 << uint32(statusClass-1)

			if (atomic.LoadUint32(config.LogBitmask))&bit > 0 {
				rw.end = time.Now()

				buf := bytes.NewBuffer(make([]byte, 0, 128))
				for _, emitter := range emitters {
					emitter(rw, r, buf)
				}
				_, _ = buf.WriteTo(config.LogWriter)
			}
		}
	})
}
<|endoftext|>