<commit_before>// Copyright (c) 2013-2014 by Michael Dvorkin. All Rights Reserved.
// Use of this source code is governed by a MIT-style license that can
// be found in the LICENSE file.

package donna

import (
	`fmt`
	`math/rand`
	`time`
)

type Globals struct {
	Log   bool // Enable logging.
	Trace bool // Trace evaluation scores.
	Fancy bool // Represent pieces as UTF-8 characters.
}

var Settings Globals

// Returns row number for the given bit index.
func Row(n int) int {
	return n >> 3 // n / 8
}

// Returns column number for the given bit index.
func Col(n int) int {
	return n & 7 // n % 8
}

// Returns row and column numbers for the given bit index.
func Coordinate(n int) (int, int) {
	return Row(n), Col(n)
}

func RelRow(square, color int) int {
	return Row(square) ^ (color * 7)
}

// Returns 0..63 square number for the given row/column coordinate.
func Square(row, column int) int {
	return (row << 3) + column
}

func Flip(color, square int) int {
	if color == White {
		return square ^ 56
	}
	return square
}

// Returns bitmask with light or dark squares set, based on color of the square.
func Same(square int) Bitmask {
	return (bit[square] & maskDark) | (bit[square] & ^maskDark)
}

func IsBetween(from, to, between int) bool {
	return ((maskStraight[from][to] | maskDiagonal[from][to]) & bit[between]) != 0
}

func Ply() int {
	return node - rootNode
}

// Integer version of math/abs.
func Abs(n int) int {
	if n < 0 {
		return -n
	}
	return n
}

func Min(x, y int) int {
	if x < y {
		return x
	}
	return y
}

func Max(x, y int) int {
	if x > y {
		return x
	}
	return y
}

// Returns, as an integer, a non-negative pseudo-random number
// in [0, limit) range. It panics if limit <= 0.
func Random(limit int) int {
	rand.Seed(time.Now().Unix())
	return rand.Intn(limit)
}

func C(color int) string {
	return [2]string{`white`, `black`}[color]
}

//
//   noWe         nort         noEa
//           +7    +8    +9
//               \  |  /
//   west    -1 <-  0 -> +1    east
//               /  |  \
//           -9    -8    -7
//   soWe         sout         soEa
//
func Rose(direction int) int {
	return [8]int{8, 9, 1, -7, -8, -9, -1, 7}[direction]
}

func Summary(metrics map[string]interface{}) {
	phase := metrics[`Phase`].(int)
	tally := metrics[`PST`].(Score)
	final := metrics[`Final`].(Score)
	units := float32(valuePawn.endgame)

	fmt.Println()
	fmt.Printf("Metric              MidGame        |        EndGame        | Blended\n")
	fmt.Printf("                W      B     W-B   |    W      B     W-B   |  (%d)  \n", phase)
	fmt.Printf("-----------------------------------+-----------------------+--------\n")
	fmt.Printf("%-12s    -      -    %5.2f  |    -      -    %5.2f  >  %5.2f\n", `PST`,
		float32(tally.midgame)/units, float32(tally.endgame)/units, float32(tally.blended(phase))/units)

	for _, tag := range([]string{`Tempo`, `Threats`, `Pawns`, `Passers`, `+Pieces`, `-Knights`, `-Bishops`, `-Rooks`, `-Queens`, `+King`, `-Cover`, `-Safety`}) {
		white := metrics[tag].(Total).white
		black := metrics[tag].(Total).black

		var score Score
		score.add(white).subtract(black)

		if tag[0:1] == `+` {
			tag = tag[1:]
		} else if tag[0:1] == `-` {
			tag = `  ` + tag[1:]
		}

		fmt.Printf("%-12s  %5.2f  %5.2f  %5.2f  |  %5.2f  %5.2f  %5.2f  >  %5.2f\n", tag,
			float32(white.midgame)/units, float32(black.midgame)/units, float32(score.midgame)/units,
			float32(white.endgame)/units, float32(black.endgame)/units, float32(score.endgame)/units,
			float32(score.blended(phase))/units)
	}
	fmt.Printf("%-12s    -      -    %5.2f  |    -      -    %5.2f  >  %5.2f\n\n", `Final Score`,
		float32(final.midgame)/units, float32(final.endgame)/units, float32(final.blended(phase))/units)
}

func SummaryAlt(metrics map[string]interface{}) {
	phase := metrics[`Phase`].(int)
	tally := metrics[`PST`].(Score)
	final := metrics[`Final`].(Score)
	units := float32(valuePawn.endgame)

	fmt.Println()
	fmt.Printf("Metric            White     |     Black     |     Total     | Blended \n")
	fmt.Printf("                mid   end   |   mid   end   |   mid   end   |  (%d)  \n", phase)
	fmt.Printf("----------------------------+---------------+---------------+---------\n")
	fmt.Printf("%-12s     -     -    |    -     -    | %5.2f  %5.2f  > %5.2f\n", `PST`,
		float32(tally.midgame)/units, float32(tally.endgame)/units,
		float32(tally.blended(phase))/units)

	for _, tag := range([]string{`Tempo`, `Threats`, `Pawns`, `Passers`, `+Pieces`, `-Knights`, `-Bishops`, `-Rooks`, `-Queens`, `+King`, `-Cover`, `-Safety`}) {
		white := metrics[tag].(Total).white
		black := metrics[tag].(Total).black

		var score Score
		score.add(white).subtract(black)

		if tag[0:1] == `+` {
			tag = tag[1:]
		} else if tag[0:1] == `-` {
			tag = `  ` + tag[1:]
		}

		fmt.Printf("%-12s  %5.2f  %5.2f  | %5.2f  %5.2f  | %5.2f  %5.2f  > %5.2f\n", tag,
			float32(white.midgame)/units, float32(white.endgame)/units,
			float32(black.midgame)/units, float32(black.endgame)/units,
			float32(score.midgame)/units, float32(score.endgame)/units,
			float32(score.blended(phase))/units)
	}
	fmt.Printf("%-12s     -     -    |    -     -    | %5.2f  %5.2f  > %5.2f\n\n", `Final Score`,
		float32(final.midgame)/units, float32(final.endgame)/units,
		float32(final.blended(phase))/units)
}

// Logging wrapper around fmt.Printf() that could be turned on as needed. Typical
// usage is Log(true); defer Log(false) in tests.
func Log(args ...interface{}) {
	switch len(args) {
	case 0:
		// Calling Log() with no arguments flips the logging setting.
		Settings.Log = !Settings.Log
		Settings.Fancy = !Settings.Fancy
	case 1:
		switch args[0].(type) {
		case bool:
			Settings.Log = args[0].(bool)
			Settings.Fancy = args[0].(bool)
		default:
			if Settings.Log {
				fmt.Println(args...)
			}
		}
	default:
		if Settings.Log {
			fmt.Printf(args[0].(string), args[1:]...)
		}
	}
}
<commit_msg>Display material balance when eveluation tracing is enabled<commit_after>// Copyright (c) 2013-2014 by Michael Dvorkin. All Rights Reserved.
// Use of this source code is governed by a MIT-style license that can
// be found in the LICENSE file.

package donna

import (
	`fmt`
	`math/rand`
	`time`
)

type Globals struct {
	Log   bool // Enable logging.
	Trace bool // Trace evaluation scores.
	Fancy bool // Represent pieces as UTF-8 characters.
}

var Settings Globals

// Returns row number for the given bit index.
func Row(n int) int {
	return n >> 3 // n / 8
}

// Returns column number for the given bit index.
func Col(n int) int {
	return n & 7 // n % 8
}

// Returns row and column numbers for the given bit index.
func Coordinate(n int) (int, int) {
	return Row(n), Col(n)
}

func RelRow(square, color int) int {
	return Row(square) ^ (color * 7)
}

// Returns 0..63 square number for the given row/column coordinate.
func Square(row, column int) int {
	return (row << 3) + column
}

func Flip(color, square int) int {
	if color == White {
		return square ^ 56
	}
	return square
}

// Returns bitmask with light or dark squares set, based on color of the square.
func Same(square int) Bitmask {
	return (bit[square] & maskDark) | (bit[square] & ^maskDark)
}

func IsBetween(from, to, between int) bool {
	return ((maskStraight[from][to] | maskDiagonal[from][to]) & bit[between]) != 0
}

func Ply() int {
	return node - rootNode
}

// Integer version of math/abs.
func Abs(n int) int {
	if n < 0 {
		return -n
	}
	return n
}

func Min(x, y int) int {
	if x < y {
		return x
	}
	return y
}

func Max(x, y int) int {
	if x > y {
		return x
	}
	return y
}

// Returns, as an integer, a non-negative pseudo-random number
// in [0, limit) range. It panics if limit <= 0.
func Random(limit int) int {
	rand.Seed(time.Now().Unix())
	return rand.Intn(limit)
}

func C(color int) string {
	return [2]string{`white`, `black`}[color]
}

//
//   noWe         nort         noEa
//           +7    +8    +9
//               \  |  /
//   west    -1 <-  0 -> +1    east
//               /  |  \
//           -9    -8    -7
//   soWe         sout         soEa
//
func Rose(direction int) int {
	return [8]int{8, 9, 1, -7, -8, -9, -1, 7}[direction]
}

func Summary(metrics map[string]interface{}) {
	phase := metrics[`Phase`].(int)
	tally := metrics[`PST`].(Score)
	material := metrics[`Material`].(Score)
	final := metrics[`Final`].(Score)
	units := float32(valuePawn.endgame)

	fmt.Println()
	fmt.Printf("Metric              MidGame        |        EndGame        | Blended\n")
	fmt.Printf("                W      B     W-B   |    W      B     W-B   |  (%d)  \n", phase)
	fmt.Printf("-----------------------------------+-----------------------+--------\n")
	fmt.Printf("%-12s    -      -    %5.2f  |    -      -    %5.2f  >  %5.2f\n", `PST`,
		float32(tally.midgame)/units, float32(tally.endgame)/units, float32(tally.blended(phase))/units)
	fmt.Printf("%-12s    -      -    %5.2f  |    -      -    %5.2f  >  %5.2f\n", `Material`,
		float32(material.midgame)/units, float32(material.endgame)/units, float32(material.blended(phase))/units)

	for _, tag := range([]string{`Tempo`, `Threats`, `Pawns`, `Passers`, `+Pieces`, `-Knights`, `-Bishops`, `-Rooks`, `-Queens`, `+King`, `-Cover`, `-Safety`}) {
		white := metrics[tag].(Total).white
		black := metrics[tag].(Total).black

		var score Score
		score.add(white).subtract(black)

		if tag[0:1] == `+` {
			tag = tag[1:]
		} else if tag[0:1] == `-` {
			tag = `  ` + tag[1:]
		}

		fmt.Printf("%-12s  %5.2f  %5.2f  %5.2f  |  %5.2f  %5.2f  %5.2f  >  %5.2f\n", tag,
			float32(white.midgame)/units, float32(black.midgame)/units, float32(score.midgame)/units,
			float32(white.endgame)/units, float32(black.endgame)/units, float32(score.endgame)/units,
			float32(score.blended(phase))/units)
	}
	fmt.Printf("%-12s    -      -    %5.2f  |    -      -    %5.2f  >  %5.2f\n\n", `Final Score`,
		float32(final.midgame)/units, float32(final.endgame)/units, float32(final.blended(phase))/units)
}

func SummaryAlt(metrics map[string]interface{}) {
	phase := metrics[`Phase`].(int)
	tally := metrics[`PST`].(Score)
	material := metrics[`Material`].(Score)
	final := metrics[`Final`].(Score)
	units := float32(valuePawn.endgame)

	fmt.Println()
	fmt.Printf("Metric            White     |     Black     |     Total     | Blended \n")
	fmt.Printf("                mid   end   |   mid   end   |   mid   end   |  (%d)  \n", phase)
	fmt.Printf("----------------------------+---------------+---------------+---------\n")
	fmt.Printf("%-12s     -     -    |    -     -    | %5.2f  %5.2f  > %5.2f\n", `PST`,
		float32(tally.midgame)/units, float32(tally.endgame)/units, float32(tally.blended(phase))/units)
	fmt.Printf("%-12s    -      -    %5.2f  |    -      -    %5.2f  >  %5.2f\n", `Material`,
		float32(material.midgame)/units, float32(material.endgame)/units, float32(material.blended(phase))/units)

	for _, tag := range([]string{`Tempo`, `Threats`, `Pawns`, `Passers`, `+Pieces`, `-Knights`, `-Bishops`, `-Rooks`, `-Queens`, `+King`, `-Cover`, `-Safety`}) {
		white := metrics[tag].(Total).white
		black := metrics[tag].(Total).black

		var score Score
		score.add(white).subtract(black)

		if tag[0:1] == `+` {
			tag = tag[1:]
		} else if tag[0:1] == `-` {
			tag = `  ` + tag[1:]
		}

		fmt.Printf("%-12s  %5.2f  %5.2f  | %5.2f  %5.2f  | %5.2f  %5.2f  > %5.2f\n", tag,
			float32(white.midgame)/units, float32(white.endgame)/units,
			float32(black.midgame)/units, float32(black.endgame)/units,
			float32(score.midgame)/units, float32(score.endgame)/units,
			float32(score.blended(phase))/units)
	}
	fmt.Printf("%-12s     -     -    |    -     -    | %5.2f  %5.2f  > %5.2f\n\n", `Final Score`,
		float32(final.midgame)/units, float32(final.endgame)/units,
		float32(final.blended(phase))/units)
}

// Logging wrapper around fmt.Printf() that could be turned on as needed. Typical
// usage is Log(true); defer Log(false) in tests.
func Log(args ...interface{}) {
	switch len(args) {
	case 0:
		// Calling Log() with no arguments flips the logging setting.
		Settings.Log = !Settings.Log
		Settings.Fancy = !Settings.Fancy
	case 1:
		switch args[0].(type) {
		case bool:
			Settings.Log = args[0].(bool)
			Settings.Fancy = args[0].(bool)
		default:
			if Settings.Log {
				fmt.Println(args...)
			}
		}
	default:
		if Settings.Log {
			fmt.Printf(args[0].(string), args[1:]...)
		}
	}
}
<|endoftext|>


<commit_before>package cfutil

import (
	"errors"
	"fmt"
	"regexp"

	cfenv "github.com/cloudfoundry-community/go-cfenv"
	vault "github.com/hashicorp/vault/api"
)

var v1Regex = regexp.MustCompile(`/v1/`)

type VaultClient struct {
	vault.Client
	Endpoint           string
	RoleID             string
	SecretID           string
	ServiceSecretPath  string
	ServiceTransitPath string
	SpaceSecretPath    string
	OrgSecretPath      string
	Secret             *vault.Secret
}

func (v *VaultClient) Login() (err error) {
	path := "auth/approle/login"
	options := map[string]interface{}{
		"role_id":   v.RoleID,
		"secret_id": v.SecretID,
	}
	v.Secret, err = v.Logical().Write(path, options)
	v.SetToken(v.Secret.Auth.ClientToken)
	return err
}

func (v *VaultClient) ReadSpaceString(path string) (string, error) {
	return v.ReadString(v.SpaceSecretPath, path)
}

func (v *VaultClient) ReadOrgString(path string) (string, error) {
	return v.ReadString(v.OrgSecretPath, path)
}

func (v *VaultClient) ReadString(prefix, path string) (string, error) {
	err := v.Login()
	if err != nil {
		return "", err
	}
	location := prefix + "/" + path
	secret, err := v.Logical().Read(location)
	if err != nil {
		return "", err
	}
	str, ok := secret.Data["value"].(string)
	if !ok || str == "" {
		return "", fmt.Errorf("Missing value on path %s", location)
	}
	return str, nil
}

func NewVaultClient(serviceName string) (*VaultClient, error) {
	appEnv, _ := Current()
	var service *cfenv.Service
	var err error
	if serviceName != "" {
		service, err = serviceByName(appEnv, serviceName)
	} else {
		service, err = serviceByTag(appEnv, "Vault")
	}
	if err != nil {
		return nil, err
	}
	if service == nil {
		return nil, errors.New("Vault service not found")
	}
	var vaultClient VaultClient

	if str, ok := service.Credentials["role_id"].(string); ok {
		vaultClient.RoleID = str
	}
	if str, ok := service.Credentials["secret_id"].(string); ok {
		vaultClient.SecretID = str
	}
	if str, ok := service.Credentials["org_secret_path"].(string); ok {
		vaultClient.OrgSecretPath = v1Regex.ReplaceAllString(str, "")
	}
	if str, ok := service.Credentials["service_secret_path"].(string); ok {
		vaultClient.ServiceSecretPath = v1Regex.ReplaceAllString(str, "")
	}
	if str, ok := service.Credentials["endpoint"].(string); ok {
		vaultClient.Endpoint = str
	}
	if str, ok := service.Credentials["space_secret_path"].(string); ok {
		vaultClient.SpaceSecretPath = v1Regex.ReplaceAllString(str, "")
	}
	if str, ok := service.Credentials["service_transit_path"].(string); ok {
		vaultClient.ServiceTransitPath = v1Regex.ReplaceAllString(str, "")
	}

	client, err := vault.NewClient(&vault.Config{
		Address: vaultClient.Endpoint,
	})
	if err != nil {
		return nil, err
	}
	vaultClient.Client = *client
	err = vaultClient.Login()
	if err != nil {
		return nil, err
	}
	return &vaultClient, vaultClient.Login()
}
<commit_msg>secret can be missing as well<commit_after>package cfutil

import (
	"errors"
	"fmt"
	"regexp"

	cfenv "github.com/cloudfoundry-community/go-cfenv"
	vault "github.com/hashicorp/vault/api"
)

var v1Regex = regexp.MustCompile(`/v1/`)

type VaultClient struct {
	vault.Client
	Endpoint           string
	RoleID             string
	SecretID           string
	ServiceSecretPath  string
	ServiceTransitPath string
	SpaceSecretPath    string
	OrgSecretPath      string
	Secret             *vault.Secret
}

func (v *VaultClient) Login() (err error) {
	path := "auth/approle/login"
	options := map[string]interface{}{
		"role_id":   v.RoleID,
		"secret_id": v.SecretID,
	}
	v.Secret, err = v.Logical().Write(path, options)
	v.SetToken(v.Secret.Auth.ClientToken)
	return err
}

func (v *VaultClient) ReadSpaceString(path string) (string, error) {
	return v.ReadString(v.SpaceSecretPath, path)
}

func (v *VaultClient) ReadOrgString(path string) (string, error) {
	return v.ReadString(v.OrgSecretPath, path)
}

func (v *VaultClient) ReadString(prefix, path string) (string, error) {
	err := v.Login()
	if err != nil {
		return "", err
	}
	location := prefix + "/" + path
	secret, err := v.Logical().Read(location)
	if err != nil {
		return "", err
	}
	if secret == nil {
		return "", fmt.Errorf("Missing value on path %s", location)
	}
	str, ok := secret.Data["value"].(string)
	if !ok || str == "" {
		return "", fmt.Errorf("Missing value on path %s", location)
	}
	return str, nil
}

func NewVaultClient(serviceName string) (*VaultClient, error) {
	appEnv, _ := Current()
	var service *cfenv.Service
	var err error
	if serviceName != "" {
		service, err = serviceByName(appEnv, serviceName)
	} else {
		service, err = serviceByTag(appEnv, "Vault")
	}
	if err != nil {
		return nil, err
	}
	if service == nil {
		return nil, errors.New("Vault service not found")
	}
	var vaultClient VaultClient

	if str, ok := service.Credentials["role_id"].(string); ok {
		vaultClient.RoleID = str
	}
	if str, ok := service.Credentials["secret_id"].(string); ok {
		vaultClient.SecretID = str
	}
	if str, ok := service.Credentials["org_secret_path"].(string); ok {
		vaultClient.OrgSecretPath = v1Regex.ReplaceAllString(str, "")
	}
	if str, ok := service.Credentials["service_secret_path"].(string); ok {
		vaultClient.ServiceSecretPath = v1Regex.ReplaceAllString(str, "")
	}
	if str, ok := service.Credentials["endpoint"].(string); ok {
		vaultClient.Endpoint = str
	}
	if str, ok := service.Credentials["space_secret_path"].(string); ok {
		vaultClient.SpaceSecretPath = v1Regex.ReplaceAllString(str, "")
	}
	if str, ok := service.Credentials["service_transit_path"].(string); ok {
		vaultClient.ServiceTransitPath = v1Regex.ReplaceAllString(str, "")
	}

	client, err := vault.NewClient(&vault.Config{
		Address: vaultClient.Endpoint,
	})
	if err != nil {
		return nil, err
	}
	vaultClient.Client = *client
	err = vaultClient.Login()
	if err != nil {
		return nil, err
	}
	return &vaultClient, vaultClient.Login()
}
<|endoftext|>


<commit_before>package venom

import (
	"strings"
	"sort"

	"github.com/spf13/pflag"
	"github.com/spf13/viper"
)

// Teach viper to search FOO_BAR for every --foo-bar key instead of
// the default FOO-BAR.
func AutomaticEnv(flags *pflag.FlagSet, viperMaybe ...*viper.Viper) {
	v := viper.GetViper()
	if len(viperMaybe) != 0 {
		v = viperMaybe[0]
	}

	replaceMap := make(map[string]string, flags.NFlag())
	flags.VisitAll(func(f *pflag.Flag) {
		name := strings.ToUpper(f.Name)
		replaceMap[name] = strings.Replace(name, "-", "_", -1)
	})

	keys := make([]string, 0, len(replaceMap))
	for k := range replaceMap {
		keys = append(keys, k)
	}

	// Reverse sort keys, this is to make sure foo-bar comes before foo. This is to prevent
	// foo being triggered when foo-bar is given to string replacer.
	sort.Sort(sort.Reverse(sort.StringSlice(keys)))

	values := make([]string, 0, 2 * len(keys))
	for _, k := range keys {
		values = append(values, k)
		values = append(values, replaceMap[k])
	}

	v.SetEnvKeyReplacer(strings.NewReplacer(values...))
	v.AutomaticEnv()
}

func TwelveFactor(name string, flags *pflag.FlagSet, viperMaybe ...*viper.Viper) error {
	v := viper.GetViper()
	if len(viperMaybe) != 0 {
		v = viperMaybe[0]
	}

	// Bind flags and configuration keys 1-to-1
	err := v.BindPFlags(flags)
	if err != nil {
		return err
	}

	// Set env prefix
	v.SetEnvPrefix(strings.ToUpper(name))

	// Patch automatic env
	AutomaticEnv(flags, v)

	return nil
}
<commit_msg>$ make fmt<commit_after>package venom

import (
	"sort"
	"strings"

	"github.com/spf13/pflag"
	"github.com/spf13/viper"
)

// Teach viper to search FOO_BAR for every --foo-bar key instead of
// the default FOO-BAR.
func AutomaticEnv(flags *pflag.FlagSet, viperMaybe ...*viper.Viper) {
	v := viper.GetViper()
	if len(viperMaybe) != 0 {
		v = viperMaybe[0]
	}

	replaceMap := make(map[string]string, flags.NFlag())
	flags.VisitAll(func(f *pflag.Flag) {
		name := strings.ToUpper(f.Name)
		replaceMap[name] = strings.Replace(name, "-", "_", -1)
	})

	keys := make([]string, 0, len(replaceMap))
	for k := range replaceMap {
		keys = append(keys, k)
	}

	// Reverse sort keys, this is to make sure foo-bar comes before foo. This is to prevent
	// foo being triggered when foo-bar is given to string replacer.
	sort.Sort(sort.Reverse(sort.StringSlice(keys)))

	values := make([]string, 0, 2*len(keys))
	for _, k := range keys {
		values = append(values, k)
		values = append(values, replaceMap[k])
	}

	v.SetEnvKeyReplacer(strings.NewReplacer(values...))
	v.AutomaticEnv()
}

func TwelveFactor(name string, flags *pflag.FlagSet, viperMaybe ...*viper.Viper) error {
	v := viper.GetViper()
	if len(viperMaybe) != 0 {
		v = viperMaybe[0]
	}

	// Bind flags and configuration keys 1-to-1
	err := v.BindPFlags(flags)
	if err != nil {
		return err
	}

	// Set env prefix
	v.SetEnvPrefix(strings.ToUpper(name))

	// Patch automatic env
	AutomaticEnv(flags, v)

	return nil
}
<|endoftext|>


<commit_before>package venom

import (
	"strings"

	"github.com/spf13/pflag"
	"github.com/spf13/viper"
)

// Teach viper to search FOO_BAR for every --foo-bar key instead of
// the default FOO-BAR.
func AutomaticEnv(flags *pflag.FlagSet, viperMaybe ...*viper.Viper) {
	v := viper.GetViper()
	if len(viperMaybe) != 0 {
		v = viperMaybe[0]
	}

	replaceMap := make([]string, 0, 2 * flags.NFlag())
	flags.VisitAll(func(f *pflag.Flag) {
		name := strings.ToUpper(f.Name)
		replaceMap = append(replaceMap, name)
		replaceMap = append(replaceMap, strings.Replace(name, "-", "_", -1))
	})
	v.SetEnvKeyReplacer(strings.NewReplacer(replaceMap...))
	v.AutomaticEnv()
}

func TwelveFactor(name string, flags *pflag.FlagSet, viperMaybe ...*viper.Viper) error {
	v := viper.GetViper()
	if len(viperMaybe) != 0 {
		v = viperMaybe[0]
	}

	// Bind flags and configuration keys 1-to-1
	err := v.BindPFlags(flags)
	if err != nil {
		return err
	}

	// Set env prefix
	v.SetEnvPrefix(strings.ToUpper(name))

	// Patch automatic env
	AutomaticEnv(flags, v)

	return nil
}
<commit_msg>Make sure to fix bug with invalid replacement<commit_after>package venom

import (
	"strings"
	"sort"

	"github.com/spf13/pflag"
	"github.com/spf13/viper"
)

// Teach viper to search FOO_BAR for every --foo-bar key instead of
// the default FOO-BAR.
func AutomaticEnv(flags *pflag.FlagSet, viperMaybe ...*viper.Viper) {
	v := viper.GetViper()
	if len(viperMaybe) != 0 {
		v = viperMaybe[0]
	}

	replaceMap := make(map[string]string, flags.NFlag())
	flags.VisitAll(func(f *pflag.Flag) {
		name := strings.ToUpper(f.Name)
		replaceMap[name] = strings.Replace(name, "-", "_", -1)
	})

	keys := make([]string, 0, len(replaceMap))
	for k := range replaceMap {
		keys = append(keys, k)
	}

	// Reverse sort keys, this is to make sure foo-bar comes before foo. This is to prevent
	// foo being triggered when foo-bar is given to string replacer.
	sort.Sort(sort.Reverse(sort.StringSlice(keys)))

	values := make([]string, 0, 2 * len(keys))
	for _, k := range keys {
		values = append(values, k)
		values = append(values, replaceMap[k])
	}

	v.SetEnvKeyReplacer(strings.NewReplacer(values...))
	v.AutomaticEnv()
}

func TwelveFactor(name string, flags *pflag.FlagSet, viperMaybe ...*viper.Viper) error {
	v := viper.GetViper()
	if len(viperMaybe) != 0 {
		v = viperMaybe[0]
	}

	// Bind flags and configuration keys 1-to-1
	err := v.BindPFlags(flags)
	if err != nil {
		return err
	}

	// Set env prefix
	v.SetEnvPrefix(strings.ToUpper(name))

	// Patch automatic env
	AutomaticEnv(flags, v)

	return nil
}
<|endoftext|>


<commit_before>// delta.go - Highlight lines with large timestamp delta.
//
// delta reads lines of text from log files or stdin, tries to find timestamps
// in those lines of text and calculates the difference of the timestamps
// between subsecquent lines. If this delta is larger then a certain limit, an
// extra line of text that visually seperates those two lines is inserted.
//
// TODO: automatic duration limit

package main

import "bufio"
import "flag"
import "fmt"
import "io"
import "os"
import "regexp"
import "time"

// Two subsequent lines with timestamp differences larger then
// timestampDifferenceLimit will get seperated.
var timestampDifferenceLimit time.Duration

// Name of the input file ("" if stdin is used).
var inputFileName string

// Holds the timestamp of the previous line.
var previousTimestamp time.Time

// Type for timestamp formats. The fields definition and example are used to
// specify a timestamp. The field compiled holds the timestamp as compiled
// regular expression.
type TimestampFormat struct {
	definition string
	example    string
	compiled   regexp.Regexp
}

// Optional custom format from the command line.
var customFormat TimestampFormat

// All specified timestamp formats.
var timestampFormats []TimestampFormat

// Prepare predefined or custom timestamp formats.
func prepareTimestampFormats() {
	if customFormat.definition == "" {
		// hh:mm:ss.mmmuuu (glog)
		timestampFormats = append(timestampFormats, TimestampFormat{
			definition: "(?P<time>[0-9]{2}:[0-9]{2}:[0-9]{2}.[0-9]{6})",
			example:    "15:04:05.000000"})
		// hh:mm:ss.mmm
		timestampFormats = append(timestampFormats, TimestampFormat{
			definition: "(?P<time>[0-9]{2}:[0-9]{2}:[0-9]{2}.[0-9]{3})",
			example:    "15:04:05.000"})
		// hh:mm:ss
		timestampFormats = append(timestampFormats, TimestampFormat{
			definition: "(?P<time>[0-9]{2}:[0-9]{2}:[0-9]{2})",
			example:    "15:04:05"})
		// hh:mm
		timestampFormats = append(timestampFormats, TimestampFormat{
			definition: "(?P<time>[0-9]{2}:[0-9]{2})",
			example:    "15:04"})
	} else {
		timestampFormats = append(timestampFormats, TimestampFormat{
			definition: customFormat.definition,
			example:    customFormat.example})
	}
	// Compile regular expressions.
	for i := 0; i < len(timestampFormats); i++ {
		compiled, err := regexp.Compile(timestampFormats[i].definition)
		if err != nil {
			panic(err)
		}
		timestampFormats[i].compiled = *compiled
	}
}

// Holds the seperator configuration.
var seperator struct {
	pattern string
	reps    int
	line    string
}

// Prepare the seperator line. We only want to do this once.
func prepareSeperator() {
	for i := 0; i < seperator.reps; i++ {
		seperator.line += seperator.pattern
	}
}

// Check if we have a large timestamp difference.
func largeTimestampDifference(t time.Time) bool {
	diff := -previousTimestamp.Sub(t)
	previousTimestamp = t
	if diff > timestampDifferenceLimit {
		return true
	}
	return false
}

// Analyze a single line.
func analyzeLine(line []byte) {
	// Check if any of the known timestamp formats fits.
	for i := 0; i < len(timestampFormats); i++ {
		regexp := timestampFormats[i].compiled
		tuple := regexp.FindSubmatchIndex(line)
		if tuple != nil {
			start := tuple[0]
			end := tuple[1]
			raw := line[start:end]
			parsed, err := time.Parse(timestampFormats[i].example, string(raw))
			if err != nil {
				continue
			}
			//fmt.Println("timestamp", parsed)
			if largeTimestampDifference(parsed) {
				fmt.Println(seperator.line)
			}
			break
		}
	}
	fmt.Println(string(line))
}

// Analyze proper.
func analyze(reader *bufio.Reader) {
	// Main loop.
	for {
		line, err := reader.ReadString('\n')
		if err == nil {
			analyzeLine([]byte(line[:len(line)-1]))
		} else if err == io.EOF {
			return
		} else {
			panic(err)
		}
	}
}

// Read from file.
func analyzeFile(filename string) {
	file, err := os.Open(filename)
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
	defer file.Close()
	reader := bufio.NewReader(file)
	analyze(reader)
}

// Read from stdin.
func analyzeStdin() {
	reader := bufio.NewReader(os.Stdin)
	analyze(reader)
}

// Print usage message.
func usage() {
	fmt.Println(`Usage: delta <[FILE] >[FILE]

	tail -f /var/log/messages | delta

delta - highlight timestamp gaps.

It reads from stdin, tries to find timestamps and calculates the timestamp
delta between subsequent lines. If this delta is larger then a certain limit,
an extra line of ASCII decoration that visually seperates those two lines is
inserted.

Options:`)
	flag.PrintDefaults()
	os.Exit(2)
}

func main() {
	// Care about command line flags.
	duration := "100ms"
	flag.Usage = usage
	flag.StringVar(&inputFileName, "f", "", "Read from this file")
	flag.StringVar(&customFormat.definition, "c", "",
		"Use a custom timestamp format instead of the predefined ones. "+
			"If used, an example has to provided with the -e switch")
	flag.StringVar(&customFormat.example, "e", "",
		"Example for the custom timestamp format")
	flag.StringVar(&duration, "d", duration,
		"Duration limit with unit suffix, e.g. 250ms, 1h45m. Valid time "+
			"units are ns, us, ms, s, m, h")
	flag.StringVar(&seperator.pattern, "p", "-",
		"Defines a custom seperator pattern")
	flag.IntVar(&seperator.reps, "r", 80,
		"Defines how often the seperator pattern will be repeated")
	flag.Parse()
	d, err := time.ParseDuration(duration)
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	} else {
		timestampDifferenceLimit = d
	}

	// Compile regular expressions and prepare seperator.
	prepareTimestampFormats()
	prepareSeperator()

	// Do work.
	if inputFileName == "" {
		analyzeStdin()
	} else {
		analyzeFile(inputFileName)
	}
}

<commit_msg>Preperations for dynamic duration limits.<commit_after>// delta.go - Highlight timestamp gaps.
//
// delta reads lines of text from log files or stdin, tries to find timestamps
// in those lines of text and calculates the difference of the timestamps
// between subsecquent lines. If this delta is larger then a certain limit, an
// extra line of text that visually seperates those two lines is inserted.
//
// TODO: automatic duration limit

package main

import "bufio"
import "flag"
import "fmt"
import "io"
import "os"
import "regexp"
import "time"

// Two subsequent lines with timestamp differences larger then
// timestampDifferenceLimit will get seperated.
var timestampDifferenceLimit time.Duration

// TODO doc
var comparisionFunction func(t time.Time) bool = largeTimestampDifference

// Name of the input file ("" if stdin is used).
var inputFileName string

// Holds the timestamp of the previous line.
var previousTimestamp time.Time

// Type for timestamp formats. The fields definition and example are used to
// specify a timestamp. The field compiled holds the timestamp as compiled
// regular expression.
type TimestampFormat struct {
	definition string
	example    string
	compiled   regexp.Regexp
}

// Optional custom format from the command line.
var customFormat TimestampFormat

// All specified timestamp formats.
var timestampFormats []TimestampFormat

// Prepare predefined or custom timestamp formats.
func prepareTimestampFormats() {
	if customFormat.definition == "" {
		// hh:mm:ss.mmmuuu (glog)
		timestampFormats = append(timestampFormats, TimestampFormat{
			definition: "(?P<time>[0-9]{2}:[0-9]{2}:[0-9]{2}.[0-9]{6})",
			example:    "15:04:05.000000"})
		// hh:mm:ss.mmm
		timestampFormats = append(timestampFormats, TimestampFormat{
			definition: "(?P<time>[0-9]{2}:[0-9]{2}:[0-9]{2}.[0-9]{3})",
			example:    "15:04:05.000"})
		// hh:mm:ss
		timestampFormats = append(timestampFormats, TimestampFormat{
			definition: "(?P<time>[0-9]{2}:[0-9]{2}:[0-9]{2})",
			example:    "15:04:05"})
		// hh:mm
		timestampFormats = append(timestampFormats, TimestampFormat{
			definition: "(?P<time>[0-9]{2}:[0-9]{2})",
			example:    "15:04"})
	} else {
		timestampFormats = append(timestampFormats, TimestampFormat{
			definition: customFormat.definition,
			example:    customFormat.example})
	}
	// Compile regular expressions.
	for i := 0; i < len(timestampFormats); i++ {
		compiled, err := regexp.Compile(timestampFormats[i].definition)
		if err != nil {
			panic(err)
		}
		timestampFormats[i].compiled = *compiled
	}
}

// Holds the seperator configuration.
var seperator struct {
	pattern string
	reps    int
	line    string
}

// Prepare the seperator line. We only want to do this once.
func prepareSeperator() {
	for i := 0; i < seperator.reps; i++ {
		seperator.line += seperator.pattern
	}
}

// Check if we have a large timestamp difference.
func largeTimestampDifference(t time.Time) bool {
	diff := -previousTimestamp.Sub(t)
	previousTimestamp = t
	if diff > timestampDifferenceLimit {
		return true
	}
	return false
}

var first int64 = 0
var mean int64

// todo: doesn't work very good yet!
func dynamicDifference(t time.Time) bool {
	if first == 0 {
		previousTimestamp = t
		first = 1
		return false
	} else if first == 1 {
		diff := -previousTimestamp.Sub(t)
		mean = int64(diff)
		first = 2
		return false
	} else if first == 2 {
		diff := -previousTimestamp.Sub(t)
		previousTimestamp = t
		eta := 0.1
		mean += int64(float64((int64(diff) - mean)) * eta)
		fmt.Printf("diff = %12d, mean = %12d |", diff/1000000, mean/1000000)
		if int64(diff) > mean {
			return false // true
		}
		return false
	}
	return false
}

// Analyze a single line.
func analyzeLine(line []byte) {
	// Check if any of the known timestamp formats fits.
	for i := 0; i < len(timestampFormats); i++ {
		regexp := timestampFormats[i].compiled
		tuple := regexp.FindSubmatchIndex(line)
		if tuple != nil {
			start := tuple[0]
			end := tuple[1]
			raw := line[start:end]
			parsed, err := time.Parse(timestampFormats[i].example, string(raw))
			if err != nil {
				continue
			}
			if comparisionFunction(parsed) {
				fmt.Println(seperator.line)
			}
			break
		}
	}
	fmt.Println(string(line))
}

// Analyze proper.
func analyze(reader *bufio.Reader) {
	// Main loop.
	for {
		line, err := reader.ReadString('\n')
		if err == nil {
			analyzeLine([]byte(line[:len(line)-1]))
		} else if err == io.EOF {
			return
		} else {
			panic(err)
		}
	}
}

// Read from file.
func analyzeFile(filename string) {
	file, err := os.Open(filename)
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
	defer file.Close()
	reader := bufio.NewReader(file)
	analyze(reader)
}

// Read from stdin.
func analyzeStdin() {
	reader := bufio.NewReader(os.Stdin)
	analyze(reader)
}

// Print usage message.
func usage() {
	fmt.Println(`Usage: delta <[FILE] >[FILE]

	tail -f /var/log/messages | delta

delta - highlight timestamp gaps.

It reads from stdin, tries to find timestamps and calculates the timestamp
delta between subsequent lines. If this delta is larger then a certain limit,
an extra line of ASCII decoration that visually seperates those two lines is
inserted.

Options:`)
	flag.PrintDefaults()
	os.Exit(2)
}

func main() {
	// Care about command line flags.
	dynamic := false
	duration := "100ms"
	flag.Usage = usage
	flag.StringVar(&inputFileName, "f", "", "Read from this file")
	flag.StringVar(&customFormat.definition, "c", "",
		"Use a custom timestamp format instead of the predefined ones. "+
			"If used, an example has to provided with the -e switch")
	flag.StringVar(&customFormat.example, "e", "",
		"Example for the custom timestamp format")
	flag.StringVar(&duration, "d", duration,
		"Duration limit with unit suffix, e.g. 250ms, 1h45m. Valid time "+
			"units are ns, us, ms, s, m, h")
	flag.StringVar(&seperator.pattern, "p", "-",
		"Defines a custom seperator pattern")
	flag.IntVar(&seperator.reps, "r", 80,
		"Defines how often the seperator pattern will be repeated")
	flag.BoolVar(&dynamic, "y", dynamic, "Use dynamic duration limit")
	flag.Parse()
	if dynamic {
		comparisionFunction = dynamicDifference
	}
	d, err := time.ParseDuration(duration)
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	} else {
		timestampDifferenceLimit = d
	}

	// Compile regular expressions and prepare seperator.
	prepareTimestampFormats()
	prepareSeperator()

	// Do work.
	if inputFileName == "" {
		analyzeStdin()
	} else {
		analyzeFile(inputFileName)
	}
}

<|endoftext|>


<commit_before>package nghttp2

import (
	"fmt"
	"github.com/bradfitz/http2"
	"github.com/bradfitz/http2/hpack"
	"io/ioutil"
	"net/http"
	"testing"
)

func TestH1H1PlainGET(t *testing.T) {
	st := newServerTester(nil, t, noopHandler)
	defer st.Close()

	res, err := st.http1(requestParam{
		name: "TestH1H1PlainGET",
	})
	if err != nil {
		t.Errorf("Error st.http1() = %v", err)
	}

	want := 200
	if got := res.status; got != want {
		t.Errorf("status = %v; want %v", got, want)
	}
}

func TestH1H1PlainGETClose(t *testing.T) {
	st := newServerTester(nil, t, noopHandler)
	defer st.Close()

	res, err := st.http1(requestParam{
		name: "TestH1H1PlainGET",
		header: []hpack.HeaderField{
			pair("Connection", "close"),
		},
	})
	if err != nil {
		t.Fatalf("Error st.http1() = %v", err)
	}

	want := 200
	if got := res.status; got != want {
		t.Errorf("status = %v; want %v", got, want)
	}
}

func TestH2H1PlainGET(t *testing.T) {
	st := newServerTester(nil, t, noopHandler)
	defer st.Close()

	res, err := st.http2(requestParam{
		name: "TestH2H1PlainGet",
	})
	if err != nil {
		t.Errorf("Error st.http2() = %v", err)
	}

	want := 200
	if res.status != want {
		t.Errorf("status = %v; want %v", res.status, want)
	}
}

func TestH2H1AddXff(t *testing.T) {
	st := newServerTester([]string{"--add-x-forwarded-for"}, t, func(w http.ResponseWriter, r *http.Request) {
		xff := r.Header.Get("X-Forwarded-For")
		want := "127.0.0.1"
		if xff != want {
			t.Errorf("X-Forwarded-For = %v; want %v", xff, want)
		}
	})
	defer st.Close()

	_, err := st.http2(requestParam{
		name: "TestH2H1AddXff",
	})
	if err != nil {
		t.Errorf("Error st.http2() = %v", err)
	}
}

func TestH2H1AddXff2(t *testing.T) {
	st := newServerTester([]string{"--add-x-forwarded-for"}, t, func(w http.ResponseWriter, r *http.Request) {
		xff := r.Header.Get("X-Forwarded-For")
		want := "host, 127.0.0.1"
		if xff != want {
			t.Errorf("X-Forwarded-For = %v; want %v", xff, want)
		}
	})
	defer st.Close()

	_, err := st.http2(requestParam{
		name: "TestH2H1AddXff2",
		header: []hpack.HeaderField{
			pair("x-forwarded-for", "host"),
		},
	})
	if err != nil {
		t.Errorf("Error st.http2() = %v", err)
	}
}

func TestH2H1StripXff(t *testing.T) {
	st := newServerTester([]string{"--strip-incoming-x-forwarded-for"}, t, func(w http.ResponseWriter, r *http.Request) {
		if xff, found := r.Header["X-Forwarded-For"]; found {
			t.Errorf("X-Forwarded-For = %v; want nothing", xff)
		}
	})
	defer st.Close()

	_, err := st.http2(requestParam{
		name: "TestH2H1StripXff1",
		header: []hpack.HeaderField{
			pair("x-forwarded-for", "host"),
		},
	})
	if err != nil {
		t.Errorf("Error st.http2() = %v", err)
	}
}

func TestH2H1StripAddXff(t *testing.T) {
	args := []string{
		"--strip-incoming-x-forwarded-for",
		"--add-x-forwarded-for",
	}
	st := newServerTester(args, t, func(w http.ResponseWriter, r *http.Request) {
		xff := r.Header.Get("X-Forwarded-For")
		want := "127.0.0.1"
		if xff != want {
			t.Errorf("X-Forwarded-For = %v; want %v", xff, want)
		}
	})
	defer st.Close()

	_, err := st.http2(requestParam{
		name: "TestH2H1StripAddXff",
		header: []hpack.HeaderField{
			pair("x-forwarded-for", "host"),
		},
	})
	if err != nil {
		t.Errorf("Error st.http2() = %v", err)
	}
}

func TestH2H1BadRequestCL(t *testing.T) {
	st := newServerTester(nil, t, noopHandler)
	defer st.Close()

	// we set content-length: 1024, but the actual request body is
	// 3 bytes.
	res, err := st.http2(requestParam{
		name:   "TestH2H1BadRequestCL",
		method: "POST",
		header: []hpack.HeaderField{
			pair("content-length", "1024"),
		},
		body: []byte("foo"),
	})
	if err != nil {
		t.Errorf("Error st.http2() = %v", err)
	}

	want := http2.ErrCodeProtocol
	if res.errCode != want {
		t.Errorf("res.errCode = %v; want %v", res.errCode, want)
	}
}

func TestH2H1BadResponseCL(t *testing.T) {
	st := newServerTester(nil, t, func(w http.ResponseWriter, r *http.Request) {
		// we set content-length: 1024, but only send 3 bytes.
		w.Header().Add("Content-Length", "1024")
		w.Write([]byte("foo"))
	})
	defer st.Close()

	res, err := st.http2(requestParam{
		name: "TestH2H1BadResponseCL",
	})
	if err != nil {
		t.Errorf("Error st.http2() = %v", err)
	}

	want := http2.ErrCodeProtocol
	if res.errCode != want {
		t.Errorf("res.errCode = %v; want %v", res.errCode, want)
	}
}

func TestH2H1LocationRewrite(t *testing.T) {
	st := newServerTester(nil, t, func(w http.ResponseWriter, r *http.Request) {
		// TODO we cannot get st.ts's port number here.. 8443
		// is just a place holder.  We ignore it on rewrite.
		w.Header().Add("Location", "http://127.0.0.1:8443/p/q?a=b#fragment")
	})
	defer st.Close()

	res, err := st.http2(requestParam{
		name: "TestH2H1LocationRewrite",
	})
	if err != nil {
		t.Errorf("Error st.http2() = %v", err)
	}

	want := fmt.Sprintf("http://127.0.0.1:%v/p/q?a=b#fragment", serverPort)
	if got := res.header.Get("Location"); got != want {
		t.Errorf("Location: %v; want %v", got, want)
	}
}

func TestH2H1ChunkedRequestBody(t *testing.T) {
	st := newServerTester(nil, t, func(w http.ResponseWriter, r *http.Request) {
		want := "[chunked]"
		if got := fmt.Sprint(r.TransferEncoding); got != want {
			t.Errorf("Transfer-Encoding: %v; want %v", got, want)
		}
		body, err := ioutil.ReadAll(r.Body)
		if err != nil {
			t.Errorf("Error reading r.body: %v", err)
		}
		want = "foo"
		if got := string(body); got != want {
			t.Errorf("body: %v; want %v", got, want)
		}
	})
	defer st.Close()

	_, err := st.http2(requestParam{
		name:   "TestH2H1ChunkedRequestBody",
		method: "POST",
		body:   []byte("foo"),
	})
	if err != nil {
		t.Errorf("Error st.http2() = %v", err)
	}
}

func TestH2H1DuplicateRequestCL(t *testing.T) {
	st := newServerTester(nil, t, noopHandler)
	defer st.Close()

	res, err := st.http2(requestParam{
		name: "TestH2H1DuplicateRequestCL",
		header: []hpack.HeaderField{
			pair("content-length", "1"),
			pair("content-length", "2"),
		},
	})
	if err != nil {
		t.Errorf("Error st.http2() = %v", err)
	}
	want := 400
	if got := res.status; got != want {
		t.Errorf("status: %v; want %v", got, want)
	}
}

func TestH2H1InvalidRequestCL(t *testing.T) {
	st := newServerTester(nil, t, noopHandler)
	defer st.Close()

	res, err := st.http2(requestParam{
		name: "TestH2H1InvalidRequestCL",
		header: []hpack.HeaderField{
			pair("content-length", ""),
		},
	})
	if err != nil {
		t.Errorf("Error st.http2() = %v", err)
	}
	want := 400
	if got := res.status; got != want {
		t.Errorf("status: %v; want %v", got, want)
	}
}

func TestH2H2DuplicateResponseCL(t *testing.T) {
	st := newServerTester([]string{"--http2-bridge"}, t, func(w http.ResponseWriter, r *http.Request) {
		w.Header().Add("content-length", "1")
		w.Header().Add("content-length", "2")
	})
	defer st.Close()

	res, err := st.http2(requestParam{
		name: "TestH2H2DuplicateResponseCL",
	})
	if err != nil {
		t.Errorf("Error st.http2() = %v", err)
	}
	want := 502
	if got := res.status; got != want {
		t.Errorf("status: %v; want %v", got, want)
	}
}

func TestH2H2InvalidResponseCL(t *testing.T) {
	st := newServerTester([]string{"--http2-bridge"}, t, func(w http.ResponseWriter, r *http.Request) {
		w.Header().Add("content-length", "")
	})
	defer st.Close()

	res, err := st.http2(requestParam{
		name: "TestH2H2InvalidResponseCL",
	})
	if err != nil {
		t.Errorf("Error st.http2() = %v", err)
	}
	want := 502
	if got := res.status; got != want {
		t.Errorf("status: %v; want %v", got, want)
	}
}
<commit_msg>integration: Fix minor typo<commit_after>package nghttp2

import (
	"fmt"
	"github.com/bradfitz/http2"
	"github.com/bradfitz/http2/hpack"
	"io/ioutil"
	"net/http"
	"testing"
)

func TestH1H1PlainGET(t *testing.T) {
	st := newServerTester(nil, t, noopHandler)
	defer st.Close()

	res, err := st.http1(requestParam{
		name: "TestH1H1PlainGET",
	})
	if err != nil {
		t.Errorf("Error st.http1() = %v", err)
	}

	want := 200
	if got := res.status; got != want {
		t.Errorf("status = %v; want %v", got, want)
	}
}

func TestH1H1PlainGETClose(t *testing.T) {
	st := newServerTester(nil, t, noopHandler)
	defer st.Close()

	res, err := st.http1(requestParam{
		name: "TestH1H1PlainGETClose",
		header: []hpack.HeaderField{
			pair("Connection", "close"),
		},
	})
	if err != nil {
		t.Fatalf("Error st.http1() = %v", err)
	}

	want := 200
	if got := res.status; got != want {
		t.Errorf("status = %v; want %v", got, want)
	}
}

func TestH2H1PlainGET(t *testing.T) {
	st := newServerTester(nil, t, noopHandler)
	defer st.Close()

	res, err := st.http2(requestParam{
		name: "TestH2H1PlainGET",
	})
	if err != nil {
		t.Errorf("Error st.http2() = %v", err)
	}

	want := 200
	if res.status != want {
		t.Errorf("status = %v; want %v", res.status, want)
	}
}

func TestH2H1AddXff(t *testing.T) {
	st := newServerTester([]string{"--add-x-forwarded-for"}, t, func(w http.ResponseWriter, r *http.Request) {
		xff := r.Header.Get("X-Forwarded-For")
		want := "127.0.0.1"
		if xff != want {
			t.Errorf("X-Forwarded-For = %v; want %v", xff, want)
		}
	})
	defer st.Close()

	_, err := st.http2(requestParam{
		name: "TestH2H1AddXff",
	})
	if err != nil {
		t.Errorf("Error st.http2() = %v", err)
	}
}

func TestH2H1AddXff2(t *testing.T) {
	st := newServerTester([]string{"--add-x-forwarded-for"}, t, func(w http.ResponseWriter, r *http.Request) {
		xff := r.Header.Get("X-Forwarded-For")
		want := "host, 127.0.0.1"
		if xff != want {
			t.Errorf("X-Forwarded-For = %v; want %v", xff, want)
		}
	})
	defer st.Close()

	_, err := st.http2(requestParam{
		name: "TestH2H1AddXff2",
		header: []hpack.HeaderField{
			pair("x-forwarded-for", "host"),
		},
	})
	if err != nil {
		t.Errorf("Error st.http2() = %v", err)
	}
}

func TestH2H1StripXff(t *testing.T) {
	st := newServerTester([]string{"--strip-incoming-x-forwarded-for"}, t, func(w http.ResponseWriter, r *http.Request) {
		if xff, found := r.Header["X-Forwarded-For"]; found {
			t.Errorf("X-Forwarded-For = %v; want nothing", xff)
		}
	})
	defer st.Close()

	_, err := st.http2(requestParam{
		name: "TestH2H1StripXff1",
		header: []hpack.HeaderField{
			pair("x-forwarded-for", "host"),
		},
	})
	if err != nil {
		t.Errorf("Error st.http2() = %v", err)
	}
}

func TestH2H1StripAddXff(t *testing.T) {
	args := []string{
		"--strip-incoming-x-forwarded-for",
		"--add-x-forwarded-for",
	}
	st := newServerTester(args, t, func(w http.ResponseWriter, r *http.Request) {
		xff := r.Header.Get("X-Forwarded-For")
		want := "127.0.0.1"
		if xff != want {
			t.Errorf("X-Forwarded-For = %v; want %v", xff, want)
		}
	})
	defer st.Close()

	_, err := st.http2(requestParam{
		name: "TestH2H1StripAddXff",
		header: []hpack.HeaderField{
			pair("x-forwarded-for", "host"),
		},
	})
	if err != nil {
		t.Errorf("Error st.http2() = %v", err)
	}
}

func TestH2H1BadRequestCL(t *testing.T) {
	st := newServerTester(nil, t, noopHandler)
	defer st.Close()

	// we set content-length: 1024, but the actual request body is
	// 3 bytes.
	res, err := st.http2(requestParam{
		name:   "TestH2H1BadRequestCL",
		method: "POST",
		header: []hpack.HeaderField{
			pair("content-length", "1024"),
		},
		body: []byte("foo"),
	})
	if err != nil {
		t.Errorf("Error st.http2() = %v", err)
	}

	want := http2.ErrCodeProtocol
	if res.errCode != want {
		t.Errorf("res.errCode = %v; want %v", res.errCode, want)
	}
}

func TestH2H1BadResponseCL(t *testing.T) {
	st := newServerTester(nil, t, func(w http.ResponseWriter, r *http.Request) {
		// we set content-length: 1024, but only send 3 bytes.
		w.Header().Add("Content-Length", "1024")
		w.Write([]byte("foo"))
	})
	defer st.Close()

	res, err := st.http2(requestParam{
		name: "TestH2H1BadResponseCL",
	})
	if err != nil {
		t.Errorf("Error st.http2() = %v", err)
	}

	want := http2.ErrCodeProtocol
	if res.errCode != want {
		t.Errorf("res.errCode = %v; want %v", res.errCode, want)
	}
}

func TestH2H1LocationRewrite(t *testing.T) {
	st := newServerTester(nil, t, func(w http.ResponseWriter, r *http.Request) {
		// TODO we cannot get st.ts's port number here.. 8443
		// is just a place holder.  We ignore it on rewrite.
		w.Header().Add("Location", "http://127.0.0.1:8443/p/q?a=b#fragment")
	})
	defer st.Close()

	res, err := st.http2(requestParam{
		name: "TestH2H1LocationRewrite",
	})
	if err != nil {
		t.Errorf("Error st.http2() = %v", err)
	}

	want := fmt.Sprintf("http://127.0.0.1:%v/p/q?a=b#fragment", serverPort)
	if got := res.header.Get("Location"); got != want {
		t.Errorf("Location: %v; want %v", got, want)
	}
}

func TestH2H1ChunkedRequestBody(t *testing.T) {
	st := newServerTester(nil, t, func(w http.ResponseWriter, r *http.Request) {
		want := "[chunked]"
		if got := fmt.Sprint(r.TransferEncoding); got != want {
			t.Errorf("Transfer-Encoding: %v; want %v", got, want)
		}
		body, err := ioutil.ReadAll(r.Body)
		if err != nil {
			t.Errorf("Error reading r.body: %v", err)
		}
		want = "foo"
		if got := string(body); got != want {
			t.Errorf("body: %v; want %v", got, want)
		}
	})
	defer st.Close()

	_, err := st.http2(requestParam{
		name:   "TestH2H1ChunkedRequestBody",
		method: "POST",
		body:   []byte("foo"),
	})
	if err != nil {
		t.Errorf("Error st.http2() = %v", err)
	}
}

func TestH2H1DuplicateRequestCL(t *testing.T) {
	st := newServerTester(nil, t, noopHandler)
	defer st.Close()

	res, err := st.http2(requestParam{
		name: "TestH2H1DuplicateRequestCL",
		header: []hpack.HeaderField{
			pair("content-length", "1"),
			pair("content-length", "2"),
		},
	})
	if err != nil {
		t.Errorf("Error st.http2() = %v", err)
	}
	want := 400
	if got := res.status; got != want {
		t.Errorf("status: %v; want %v", got, want)
	}
}

func TestH2H1InvalidRequestCL(t *testing.T) {
	st := newServerTester(nil, t, noopHandler)
	defer st.Close()

	res, err := st.http2(requestParam{
		name: "TestH2H1InvalidRequestCL",
		header: []hpack.HeaderField{
			pair("content-length", ""),
		},
	})
	if err != nil {
		t.Errorf("Error st.http2() = %v", err)
	}
	want := 400
	if got := res.status; got != want {
		t.Errorf("status: %v; want %v", got, want)
	}
}

func TestH2H2DuplicateResponseCL(t *testing.T) {
	st := newServerTester([]string{"--http2-bridge"}, t, func(w http.ResponseWriter, r *http.Request) {
		w.Header().Add("content-length", "1")
		w.Header().Add("content-length", "2")
	})
	defer st.Close()

	res, err := st.http2(requestParam{
		name: "TestH2H2DuplicateResponseCL",
	})
	if err != nil {
		t.Errorf("Error st.http2() = %v", err)
	}
	want := 502
	if got := res.status; got != want {
		t.Errorf("status: %v; want %v", got, want)
	}
}

func TestH2H2InvalidResponseCL(t *testing.T) {
	st := newServerTester([]string{"--http2-bridge"}, t, func(w http.ResponseWriter, r *http.Request) {
		w.Header().Add("content-length", "")
	})
	defer st.Close()

	res, err := st.http2(requestParam{
		name: "TestH2H2InvalidResponseCL",
	})
	if err != nil {
		t.Errorf("Error st.http2() = %v", err)
	}
	want := 502
	if got := res.status; got != want {
		t.Errorf("status: %v; want %v", got, want)
	}
}
<|endoftext|>


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

package main

import (
	"fmt"
	"os"
	"testing"

	"github.com/coreos/rkt/tests/testutils"
)

// TestDNS is checking how rkt fills /etc/resolv.conf
func TestDNS(t *testing.T) {
	imageFile := patchTestACI("rkt-inspect-exit.aci", "--exec=/inspect --print-msg=Hello --read-file")
	defer os.Remove(imageFile)
	ctx := testutils.NewRktRunCtx()
	defer ctx.Cleanup()

	for _, tt := range []struct {
		paramDNS     string
		expectedLine string
	}{
		{
			paramDNS:     "",
			expectedLine: "Cannot read file",
		},
		{
			paramDNS:     "--dns=8.8.4.4",
			expectedLine: "nameserver 8.8.4.4",
		},
		{
			paramDNS:     "--dns=8.8.8.8 --dns=8.8.4.4",
			expectedLine: "nameserver 8.8.8.8",
		},
		{
			paramDNS:     "--dns=8.8.8.8 --dns=8.8.4.4 --dns-search=search.com --dns-opt=debug",
			expectedLine: "nameserver 8.8.4.4",
		},
		{
			paramDNS:     "--dns-search=foo.com --dns-search=bar.com",
			expectedLine: "search foo.com bar.com",
		},
		{
			paramDNS:     "--dns-opt=debug --dns-opt=use-vc --dns-opt=rotate",
			expectedLine: "options debug use-vc rotate",
		},
	} {

		rktCmd := fmt.Sprintf(`%s --insecure-options=image run --set-env=FILE=/etc/resolv.conf %s %s`,
			ctx.Cmd(), tt.paramDNS, imageFile)
		t.Logf("%s\n", rktCmd)
		runRktAndCheckOutput(t, rktCmd, tt.expectedLine, false)
	}
}
<commit_msg>tests: exit code: fix TestDNS<commit_after>// Copyright 2016 The rkt Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"fmt"
	"os"
	"testing"

	"github.com/coreos/rkt/tests/testutils"
)

// TestDNS is checking how rkt fills /etc/resolv.conf
func TestDNS(t *testing.T) {
	imageFile := patchTestACI("rkt-inspect-exit.aci", "--exec=/inspect --print-msg=Hello --read-file")
	defer os.Remove(imageFile)
	ctx := testutils.NewRktRunCtx()
	defer ctx.Cleanup()

	for _, tt := range []struct {
		paramDNS      string
		expectedLine  string
		expectedError bool
	}{
		{
			paramDNS:      "",
			expectedLine:  "Cannot read file",
			expectedError: true,
		},
		{
			paramDNS:      "--dns=8.8.4.4",
			expectedLine:  "nameserver 8.8.4.4",
			expectedError: false,
		},
		{
			paramDNS:      "--dns=8.8.8.8 --dns=8.8.4.4",
			expectedLine:  "nameserver 8.8.8.8",
			expectedError: false,
		},
		{
			paramDNS:      "--dns=8.8.8.8 --dns=8.8.4.4 --dns-search=search.com --dns-opt=debug",
			expectedLine:  "nameserver 8.8.4.4",
			expectedError: false,
		},
		{
			paramDNS:      "--dns-search=foo.com --dns-search=bar.com",
			expectedLine:  "search foo.com bar.com",
			expectedError: false,
		},
		{
			paramDNS:      "--dns-opt=debug --dns-opt=use-vc --dns-opt=rotate",
			expectedLine:  "options debug use-vc rotate",
			expectedError: false,
		},
	} {

		rktCmd := fmt.Sprintf(`%s --insecure-options=image run --set-env=FILE=/etc/resolv.conf %s %s`,
			ctx.Cmd(), tt.paramDNS, imageFile)
		t.Logf("%s\n", rktCmd)
		runRktAndCheckOutput(t, rktCmd, tt.expectedLine, tt.expectedError)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/xml"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"os"
	"os/exec"
	"path"
	"text/tabwriter"
	"time"

	"code.google.com/p/go-uuid/uuid"
	"github.com/codegangsta/cli"
	"github.com/coreos/go-omaha/omaha"

	"github.com/coreos-inc/updatectl/client/update/v1"
)

func WatchCommands() []cli.Command {
	return []cli.Command{
		{
			Name:        "watch",
			Usage:       "watch [OPTION]... <appID> <groupID> <clientID> <cmd> <args>",
			Description: `Watch for app versions and exec a script`,
			Action:      handle(watch),
			Flags: []cli.Flag{
				cli.IntFlag{"interval, i", 1, "Update polling interval"},
				cli.StringFlag{"version, v", "0.0.0", "Starting version number"},
			},
		},
	}
}

func fetchUpdateCheck(server string, appID string, groupID string, clientID string, version string, debug bool) (*omaha.UpdateCheck, error) {
	client := &http.Client{}

	// TODO: Fill out the OS field correctly based on /etc/os-release
	request := omaha.NewRequest("lsb", "CoreOS", "", "")
	app := request.AddApp(fmt.Sprintf("{%s}", appID), version)
	app.AddUpdateCheck()
	app.MachineID = clientID
	app.BootId = uuid.New()
	app.Track = groupID

	event := app.AddEvent()
	event.Type = "1"
	event.Result = "0"

	raw, err := xml.MarshalIndent(request, "", " ")
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		return nil, err
	}

	if debug {
		fmt.Fprintf(os.Stderr, "Request: %s%s\n", xml.Header, raw)
	}

	resp, err := client.Post(server+"/v1/update/", "text/xml", bytes.NewReader(raw))
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		return nil, err
	}

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		return nil, err
	}

	if debug {
		fmt.Fprintf(os.Stderr, "Response: %s%s\n", xml.Header, string(body))
	}

	oresp := &omaha.Response{}
	err = xml.Unmarshal(body, oresp)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		return nil, err
	}

	return oresp.Apps[0].UpdateCheck, nil
}

func prepareEnvironment(appID string, version string, oldVersion string, updateCheck *omaha.UpdateCheck) []string {
	env := os.Environ()
	env = append(env, "UPDATE_SERVICE_VERSION="+version)
	if oldVersion != "" {
		env = append(env, "UPDATE_SERVICE_OLD_VERSION="+oldVersion)
	}
	env = append(env, "UPDATE_SERVICE_APP_ID="+appID)

	url, err := url.Parse(updateCheck.Urls.Urls[0].CodeBase)
	if err != nil {
		fmt.Fprintf(os.Stderr, err.Error())
		os.Exit(1)
	}

	url.Path = path.Join(url.Path, updateCheck.Manifest.Packages.Packages[0].Name)
	env = append(env, "UPDATE_SERVICE_URL="+url.String())
	return env
}

func runCmd(cmdName string, args []string, appID string, version string, oldVersion string, updateCheck *omaha.UpdateCheck) {
	cmd := exec.Command(cmdName, args...)
	cmd.Env = prepareEnvironment(appID, version, oldVersion, updateCheck)

	stdout, err := cmd.StdoutPipe()
	if err != nil {
		fmt.Fprintf(os.Stderr, err.Error())
		os.Exit(1)
	}
	stderr, err := cmd.StderrPipe()
	if err != nil {
		fmt.Fprintf(os.Stderr, err.Error())
		os.Exit(1)
	}
	err = cmd.Start()
	if err != nil {
		fmt.Fprintf(os.Stderr, err.Error())
		os.Exit(1)
	}
	go io.Copy(os.Stdout, stdout)
	go io.Copy(os.Stderr, stderr)
	cmd.Wait()
}

func watch(c *cli.Context, service *update.Service, out *tabwriter.Writer) {
	tick := time.NewTicker(time.Second * time.Duration(c.Int("interval")))
	server := c.GlobalString("server")
	debug := c.GlobalBool("debug")
	version := c.String("version")
	args := c.Args()

	if len(args) < 4 {
		log.Fatalf("appID, groupID and clientID required")
	}

	appID := args[0]
	groupID := args[1]
	clientID := args[2]

	// initial check
	updateCheck, err := fetchUpdateCheck(server, appID, groupID, clientID, version, debug)
	if err != nil {
		fmt.Fprintf(os.Stderr, err.Error())
		os.Exit(1)
	}
	runCmd(args[3], args[4:], appID, version, "", updateCheck)

	for {
		select {
		case <-tick.C:

			updateCheck, err := fetchUpdateCheck(server, appID, groupID, clientID, version, debug)
			if err != nil {
				log.Printf("warning: update check failed (%v)\n", err)
				continue
			}

			if updateCheck.Status == "noupdate" {
				continue
			} else if updateCheck.Status == "error-version" {
				continue
			}

			newVersion := updateCheck.Manifest.Version

			if newVersion != version {
				runCmd(args[3], args[4:], appID, newVersion, version, updateCheck)
			}
			version = newVersion
		}
	}

	tick.Stop()
}
<commit_msg>fix(watch): don't export UPDATE_SERVICE_URL if there is no update<commit_after>package main

import (
	"bytes"
	"encoding/xml"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"os"
	"os/exec"
	"path"
	"text/tabwriter"
	"time"

	"code.google.com/p/go-uuid/uuid"
	"github.com/codegangsta/cli"
	"github.com/coreos/go-omaha/omaha"

	"github.com/coreos-inc/updatectl/client/update/v1"
)

func WatchCommands() []cli.Command {
	return []cli.Command{
		{
			Name:        "watch",
			Usage:       "watch [OPTION]... <appID> <groupID> <clientID> <cmd> <args>",
			Description: `Watch for app versions and exec a script`,
			Action:      handle(watch),
			Flags: []cli.Flag{
				cli.IntFlag{"interval, i", 1, "Update polling interval"},
				cli.StringFlag{"version, v", "0.0.0", "Starting version number"},
			},
		},
	}
}

func fetchUpdateCheck(server string, appID string, groupID string, clientID string, version string, debug bool) (*omaha.UpdateCheck, error) {
	client := &http.Client{}

	// TODO: Fill out the OS field correctly based on /etc/os-release
	request := omaha.NewRequest("lsb", "CoreOS", "", "")
	app := request.AddApp(fmt.Sprintf("{%s}", appID), version)
	app.AddUpdateCheck()
	app.MachineID = clientID
	app.BootId = uuid.New()
	app.Track = groupID

	event := app.AddEvent()
	event.Type = "1"
	event.Result = "0"

	raw, err := xml.MarshalIndent(request, "", " ")
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		return nil, err
	}

	if debug {
		fmt.Fprintf(os.Stderr, "Request: %s%s\n", xml.Header, raw)
	}

	resp, err := client.Post(server+"/v1/update/", "text/xml", bytes.NewReader(raw))
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		return nil, err
	}

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		return nil, err
	}

	if debug {
		fmt.Fprintf(os.Stderr, "Response: %s%s\n", xml.Header, string(body))
	}

	oresp := &omaha.Response{}
	err = xml.Unmarshal(body, oresp)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		return nil, err
	}

	return oresp.Apps[0].UpdateCheck, nil
}

func prepareEnvironment(appID string, version string, oldVersion string, updateCheck *omaha.UpdateCheck) []string {
	env := os.Environ()
	env = append(env, "UPDATE_SERVICE_VERSION="+version)
	if oldVersion != "" {
		env = append(env, "UPDATE_SERVICE_OLD_VERSION="+oldVersion)
	}
	env = append(env, "UPDATE_SERVICE_APP_ID="+appID)

	if updateCheck.Status == "ok" {
		url, err := url.Parse(updateCheck.Urls.Urls[0].CodeBase)
		if err != nil {
			fmt.Fprintf(os.Stderr, err.Error())
			os.Exit(1)
		}

		url.Path = path.Join(url.Path, updateCheck.Manifest.Packages.Packages[0].Name)
		env = append(env, "UPDATE_SERVICE_URL="+url.String())
	}
	return env
}

func runCmd(cmdName string, args []string, appID string, version string, oldVersion string, updateCheck *omaha.UpdateCheck) {
	cmd := exec.Command(cmdName, args...)
	cmd.Env = prepareEnvironment(appID, version, oldVersion, updateCheck)

	stdout, err := cmd.StdoutPipe()
	if err != nil {
		fmt.Fprintf(os.Stderr, err.Error())
		os.Exit(1)
	}
	stderr, err := cmd.StderrPipe()
	if err != nil {
		fmt.Fprintf(os.Stderr, err.Error())
		os.Exit(1)
	}
	err = cmd.Start()
	if err != nil {
		fmt.Fprintf(os.Stderr, err.Error())
		os.Exit(1)
	}
	go io.Copy(os.Stdout, stdout)
	go io.Copy(os.Stderr, stderr)
	cmd.Wait()
}

func watch(c *cli.Context, service *update.Service, out *tabwriter.Writer) {
	tick := time.NewTicker(time.Second * time.Duration(c.Int("interval")))
	server := c.GlobalString("server")
	debug := c.GlobalBool("debug")
	version := c.String("version")
	args := c.Args()

	if len(args) < 4 {
		log.Fatalf("appID, groupID and clientID required")
	}

	appID := args[0]
	groupID := args[1]
	clientID := args[2]

	// initial check
	updateCheck, err := fetchUpdateCheck(server, appID, groupID, clientID, version, debug)
	if err != nil {
		fmt.Fprintf(os.Stderr, err.Error())
		os.Exit(1)
	}
	runCmd(args[3], args[4:], appID, version, "", updateCheck)

	for {
		select {
		case <-tick.C:

			updateCheck, err := fetchUpdateCheck(server, appID, groupID, clientID, version, debug)
			if err != nil {
				log.Printf("warning: update check failed (%v)\n", err)
				continue
			}

			if updateCheck.Status == "noupdate" {
				continue
			} else if updateCheck.Status == "error-version" {
				continue
			}

			newVersion := updateCheck.Manifest.Version

			if newVersion != version {
				runCmd(args[3], args[4:], appID, newVersion, version, updateCheck)
			}
			version = newVersion
		}
	}

	tick.Stop()
}
<|endoftext|>


<commit_before>package epub

import (
	"archive/zip"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
)

const (
	containerFilename     = "container.xml"
	containerFileTemplate = `<?xml version="1.0" encoding="UTF-8"?>
<container version="1.0" xmlns="urn:oasis:names:tc:opendocument:xmlns:container">
  <rootfiles>
    <rootfile full-path="%s/%s" media-type="application/oebps-package+xml" />
  </rootfiles>
</container>
`
	contentFolderName = "EPUB"
	// Permissions for any new directories we create
	dirPermissions = 0755
	// Permissions for any new files we create
	filePermissions   = 0644
	imageFolderName   = "img"
	mediaTypeNcx      = "application/x-dtbncx+xml"
	mediaTypeEpub     = "application/epub+zip"
	mediaTypeXhtml    = "application/xhtml+xml"
	metaInfFolderName = "META-INF"
	mimetypeFilename  = "mimetype"
	pkgFilename       = "package.opf"
	sectionFileFormat = "section%04d.xhtml"
	tempDirPrefix     = "go-epub"
	xhtmlFolderName   = "xhtml"
)

// Write writes the EPUB file. The destination path must be the full path to
// the resulting file, including filename and extension.
func (e *Epub) Write(destFilePath string) error {
	tempDir, err := ioutil.TempDir("", tempDirPrefix)
	defer os.Remove(tempDir)
	if err != nil {
		log.Fatalf("os.Remove error: %s", err)
	}

	// Must be called first
	err = createEpubFolders(tempDir)
	if err != nil {
		return err
	}

	// Must be called after:
	// createEpubFolders()
	err = e.writeImages(tempDir)
	if err != nil {
		return err
	}

	// Must be called after:
	// createEpubFolders()
	err = writeMimetype(tempDir)
	if err != nil {
		return err
	}

	// Must be called after:
	// createEpubFolders()
	err = e.writeSections(tempDir)
	if err != nil {
		return err
	}

	// Must be called after:
	// createEpubFolders()
	// writeSections()
	err = e.writeToc(tempDir)
	if err != nil {
		return err
	}

	// Must be called after:
	// createEpubFolders()
	err = writeContainerFile(tempDir)
	if err != nil {
		return err
	}

	// Must be called after:
	// createEpubFolders()
	// writeImages()
	// writeSections()
	// writeToc()
	err = e.writePackageFile(tempDir)
	if err != nil {
		return err
	}

	// Must be called last
	err = e.writeEpub(tempDir, destFilePath)
	if err != nil {
		return err
	}

	// TODO

	//	output, err := xml.MarshalIndent(e.toc.navDoc.xml, "", "  ")
	//	output = append([]byte(xhtmlDoctype), output...)

	//	output, err := xml.MarshalIndent(e.pkg.xml, "", "  ")

	//  output, err := xml.MarshalIndent(e.toc.ncxXML, "", "  ")
	//	output = append([]byte(xml.Header), output...)
	//	fmt.Println(string(output))

	return nil
}

// Create the EPUB folder structure in a temp directory
func createEpubFolders(tempDir string) error {
	if err := os.Mkdir(
		filepath.Join(
			tempDir,
			contentFolderName,
		),
		dirPermissions); err != nil {
		return err
	}

	if err := os.Mkdir(
		filepath.Join(
			tempDir,
			contentFolderName,
			xhtmlFolderName,
		),
		dirPermissions); err != nil {
		return err
	}

	if err := os.Mkdir(
		filepath.Join(
			tempDir,
			metaInfFolderName,
		),
		dirPermissions); err != nil {
		return err
	}

	return nil
}

// Write the contatiner file (container.xml), which mostly just points to the
// package file (package.opf)
//
// Sample: https://github.com/bmaupin/epub-samples/blob/master/minimal-v32/META-INF/container.xml
// Spec: http://www.idpf.org/epub/301/spec/epub-ocf.html#sec-container-metainf-container.xml
func writeContainerFile(tempDir string) error {
	containerFilePath := filepath.Join(tempDir, metaInfFolderName, containerFilename)
	if err := ioutil.WriteFile(
		containerFilePath,
		[]byte(
			fmt.Sprintf(
				containerFileTemplate,
				contentFolderName,
				pkgFilename,
			),
		),
		filePermissions,
	); err != nil {
		return err
	}

	return nil
}

// Write the EPUB file itself by zipping up everything from a temp directory
func (e *Epub) writeEpub(tempDir string, destFilePath string) error {
	f, err := os.Create(destFilePath)
	if err != nil {
		log.Fatalf("os.Create error: %s", err)
	}
	defer func() {
		if err := f.Close(); err != nil {
			log.Fatalf("os.File.Close error: %s", err)
		}
	}()

	z := zip.NewWriter(f)
	defer func() {
		if err := z.Close(); err != nil {
			log.Fatalf("zip.Writer.Close error: %s", err)
		}
	}()

	skipMimetypeFile := false

	var addFileToZip = func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}

		// Get the path of the file relative to the folder we're zipping
		relativePath, err := filepath.Rel(tempDir, path)
		if err != nil {
			log.Fatalf("filepath.Rel error: %s", err)
		}

		// Only include regular files, not directories
		if !info.Mode().IsRegular() {
			return nil
		}

		var w io.Writer
		if path == filepath.Join(tempDir, mimetypeFilename) {
			// Skip the mimetype file if it's already been written
			if skipMimetypeFile == true {
				return nil
			}
			// The mimetype file must be uncompressed according to the EPUB spec
			w, err = z.CreateHeader(&zip.FileHeader{
				Name:   relativePath,
				Method: zip.Store,
			})
		} else {
			w, err = z.Create(relativePath)
		}
		if err != nil {
			log.Fatalf("zip.Writer.Create error: %s", err)
		}

		r, err := os.Open(path)
		if err != nil {
			log.Fatalf("os.Open error: %s", err)
		}
		defer func() {
			if err := r.Close(); err != nil {
				log.Fatalf("os.File.Close error: %s", err)
			}
		}()

		_, err = io.Copy(w, r)
		if err != nil {
			log.Fatalf("io.Copy error: %s", err)
		}

		return nil
	}

	// Add the mimetype file first
	mimetypeFilePath := filepath.Join(tempDir, mimetypeFilename)
	mimetypeInfo, err := os.Lstat(mimetypeFilePath)
	if err != nil {
		log.Fatalf("os.Lstat error: %s", err)
	}
	addFileToZip(mimetypeFilePath, mimetypeInfo, nil)

	skipMimetypeFile = true

	err = filepath.Walk(tempDir, addFileToZip)
	if err != nil {
		log.Fatalf("os.Lstat error: %s", err)
	}

	return nil
}

// Get images from their source and save them in the temporary directory
func (e *Epub) writeImages(tempDir string) error {
	imageFolderPath := filepath.Join(tempDir, contentFolderName, imageFolderName)
	if err := os.Mkdir(imageFolderPath, dirPermissions); err != nil {
		return err
	}

	for imageFilename, imageSource := range e.images {
		u, err := url.Parse(imageSource)
		if err != nil {
			return err
		}

		var r io.ReadCloser
		var resp *http.Response
		// If it's a URL
		if u.Scheme == "http" || u.Scheme == "https" {
			resp, err = http.Get(imageSource)
			r = resp.Body

			// Otherwise, assume it's a local file
		} else {
			r, err = os.Open(imageSource)
		}
		if err != nil {
			return err
		}
		defer func() {
			err = r.Close()
		}()

		imageFilePath := filepath.Join(
			imageFolderPath,
			imageFilename,
		)

		w, err := os.Create(imageFilePath)
		if err != nil {
			return err
		}
		defer func() {
			err = w.Close()
		}()

		_, err = io.Copy(w, r)
		if err != nil {
			return err
		}
	}

	return nil
}

// Write the mimetype file
//
// Sample: https://github.com/bmaupin/epub-samples/blob/master/minimal-v32/mimetype
// Spec: http://www.idpf.org/epub/301/spec/epub-ocf.html#sec-zip-container-mime
func writeMimetype(tempDir string) error {
	mimetypeFilePath := filepath.Join(tempDir, mimetypeFilename)

	if err := ioutil.WriteFile(mimetypeFilePath, []byte(mediaTypeEpub), filePermissions); err != nil {
		return err
	}

	return nil
}

func (e *Epub) writePackageFile(tempDir string) error {
	err := e.pkg.write(tempDir)
	if err != nil {
		return err
	}

	return nil
}

// Write the section files to the temporary directory and add the sections to
// the TOC and package files
func (e *Epub) writeSections(tempDir string) error {
	for i, section := range e.sections {
		sectionIndex := i + 1
		sectionFilename := fmt.Sprintf(sectionFileFormat, sectionIndex)
		sectionFilePath := filepath.Join(tempDir, contentFolderName, xhtmlFolderName, sectionFilename)
		section.write(sectionFilePath)

		relativePath := filepath.Join(xhtmlFolderName, sectionFilename)
		e.toc.addSection(sectionIndex, section.Title(), relativePath)
		e.pkg.addToManifest(sectionFilename, relativePath, mediaTypeXhtml, "")
		e.pkg.addToSpine(sectionFilename)
	}

	return nil
}

// Write the TOC file to the temporary directory and add the TOC entries to the
// package file
func (e *Epub) writeToc(tempDir string) error {
	e.pkg.addToManifest(tocNavItemID, tocNavFilename, mediaTypeXhtml, tocNavItemProperties)
	e.pkg.addToManifest(tocNcxItemID, tocNcxFilename, mediaTypeNcx, "")

	err := e.toc.write(tempDir)
	if err != nil {
		return err
	}

	return nil
}
<commit_msg>Remove temporary directory after writing EPUB<commit_after>package epub

import (
	"archive/zip"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
)

const (
	containerFilename     = "container.xml"
	containerFileTemplate = `<?xml version="1.0" encoding="UTF-8"?>
<container version="1.0" xmlns="urn:oasis:names:tc:opendocument:xmlns:container">
  <rootfiles>
    <rootfile full-path="%s/%s" media-type="application/oebps-package+xml" />
  </rootfiles>
</container>
`
	contentFolderName = "EPUB"
	// Permissions for any new directories we create
	dirPermissions = 0755
	// Permissions for any new files we create
	filePermissions   = 0644
	imageFolderName   = "img"
	mediaTypeNcx      = "application/x-dtbncx+xml"
	mediaTypeEpub     = "application/epub+zip"
	mediaTypeXhtml    = "application/xhtml+xml"
	metaInfFolderName = "META-INF"
	mimetypeFilename  = "mimetype"
	pkgFilename       = "package.opf"
	sectionFileFormat = "section%04d.xhtml"
	tempDirPrefix     = "go-epub"
	xhtmlFolderName   = "xhtml"
)

// Write writes the EPUB file. The destination path must be the full path to
// the resulting file, including filename and extension.
func (e *Epub) Write(destFilePath string) error {
	tempDir, err := ioutil.TempDir("", tempDirPrefix)
	defer os.Remove(tempDir)
	if err != nil {
		log.Fatalf("os.Remove error: %s", err)
	}

	// Must be called first
	err = createEpubFolders(tempDir)
	if err != nil {
		return err
	}

	// Must be called after:
	// createEpubFolders()
	err = e.writeImages(tempDir)
	if err != nil {
		return err
	}

	// Must be called after:
	// createEpubFolders()
	err = writeMimetype(tempDir)
	if err != nil {
		return err
	}

	// Must be called after:
	// createEpubFolders()
	err = e.writeSections(tempDir)
	if err != nil {
		return err
	}

	// Must be called after:
	// createEpubFolders()
	// writeSections()
	err = e.writeToc(tempDir)
	if err != nil {
		return err
	}

	// Must be called after:
	// createEpubFolders()
	err = writeContainerFile(tempDir)
	if err != nil {
		return err
	}

	// Must be called after:
	// createEpubFolders()
	// writeImages()
	// writeSections()
	// writeToc()
	err = e.writePackageFile(tempDir)
	if err != nil {
		return err
	}

	// Must be called last
	err = e.writeEpub(tempDir, destFilePath)
	if err != nil {
		return err
	}

	// TODO

	//	output, err := xml.MarshalIndent(e.toc.navDoc.xml, "", "  ")
	//	output = append([]byte(xhtmlDoctype), output...)

	//	output, err := xml.MarshalIndent(e.pkg.xml, "", "  ")

	//  output, err := xml.MarshalIndent(e.toc.ncxXML, "", "  ")
	//	output = append([]byte(xml.Header), output...)
	//	fmt.Println(string(output))

	os.RemoveAll(tempDir)

	return nil
}

// Create the EPUB folder structure in a temp directory
func createEpubFolders(tempDir string) error {
	if err := os.Mkdir(
		filepath.Join(
			tempDir,
			contentFolderName,
		),
		dirPermissions); err != nil {
		return err
	}

	if err := os.Mkdir(
		filepath.Join(
			tempDir,
			contentFolderName,
			xhtmlFolderName,
		),
		dirPermissions); err != nil {
		return err
	}

	if err := os.Mkdir(
		filepath.Join(
			tempDir,
			metaInfFolderName,
		),
		dirPermissions); err != nil {
		return err
	}

	return nil
}

// Write the contatiner file (container.xml), which mostly just points to the
// package file (package.opf)
//
// Sample: https://github.com/bmaupin/epub-samples/blob/master/minimal-v32/META-INF/container.xml
// Spec: http://www.idpf.org/epub/301/spec/epub-ocf.html#sec-container-metainf-container.xml
func writeContainerFile(tempDir string) error {
	containerFilePath := filepath.Join(tempDir, metaInfFolderName, containerFilename)
	if err := ioutil.WriteFile(
		containerFilePath,
		[]byte(
			fmt.Sprintf(
				containerFileTemplate,
				contentFolderName,
				pkgFilename,
			),
		),
		filePermissions,
	); err != nil {
		return err
	}

	return nil
}

// Write the EPUB file itself by zipping up everything from a temp directory
func (e *Epub) writeEpub(tempDir string, destFilePath string) error {
	f, err := os.Create(destFilePath)
	if err != nil {
		log.Fatalf("os.Create error: %s", err)
	}
	defer func() {
		if err := f.Close(); err != nil {
			log.Fatalf("os.File.Close error: %s", err)
		}
	}()

	z := zip.NewWriter(f)
	defer func() {
		if err := z.Close(); err != nil {
			log.Fatalf("zip.Writer.Close error: %s", err)
		}
	}()

	skipMimetypeFile := false

	var addFileToZip = func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}

		// Get the path of the file relative to the folder we're zipping
		relativePath, err := filepath.Rel(tempDir, path)
		if err != nil {
			log.Fatalf("filepath.Rel error: %s", err)
		}

		// Only include regular files, not directories
		if !info.Mode().IsRegular() {
			return nil
		}

		var w io.Writer
		if path == filepath.Join(tempDir, mimetypeFilename) {
			// Skip the mimetype file if it's already been written
			if skipMimetypeFile == true {
				return nil
			}
			// The mimetype file must be uncompressed according to the EPUB spec
			w, err = z.CreateHeader(&zip.FileHeader{
				Name:   relativePath,
				Method: zip.Store,
			})
		} else {
			w, err = z.Create(relativePath)
		}
		if err != nil {
			log.Fatalf("zip.Writer.Create error: %s", err)
		}

		r, err := os.Open(path)
		if err != nil {
			log.Fatalf("os.Open error: %s", err)
		}
		defer func() {
			if err := r.Close(); err != nil {
				log.Fatalf("os.File.Close error: %s", err)
			}
		}()

		_, err = io.Copy(w, r)
		if err != nil {
			log.Fatalf("io.Copy error: %s", err)
		}

		return nil
	}

	// Add the mimetype file first
	mimetypeFilePath := filepath.Join(tempDir, mimetypeFilename)
	mimetypeInfo, err := os.Lstat(mimetypeFilePath)
	if err != nil {
		log.Fatalf("os.Lstat error: %s", err)
	}
	addFileToZip(mimetypeFilePath, mimetypeInfo, nil)

	skipMimetypeFile = true

	err = filepath.Walk(tempDir, addFileToZip)
	if err != nil {
		log.Fatalf("os.Lstat error: %s", err)
	}

	return nil
}

// Get images from their source and save them in the temporary directory
func (e *Epub) writeImages(tempDir string) error {
	imageFolderPath := filepath.Join(tempDir, contentFolderName, imageFolderName)
	if err := os.Mkdir(imageFolderPath, dirPermissions); err != nil {
		return err
	}

	for imageFilename, imageSource := range e.images {
		u, err := url.Parse(imageSource)
		if err != nil {
			return err
		}

		var r io.ReadCloser
		var resp *http.Response
		// If it's a URL
		if u.Scheme == "http" || u.Scheme == "https" {
			resp, err = http.Get(imageSource)
			r = resp.Body

			// Otherwise, assume it's a local file
		} else {
			r, err = os.Open(imageSource)
		}
		if err != nil {
			return err
		}
		defer func() {
			err = r.Close()
		}()

		imageFilePath := filepath.Join(
			imageFolderPath,
			imageFilename,
		)

		w, err := os.Create(imageFilePath)
		if err != nil {
			return err
		}
		defer func() {
			err = w.Close()
		}()

		_, err = io.Copy(w, r)
		if err != nil {
			return err
		}
	}

	return nil
}

// Write the mimetype file
//
// Sample: https://github.com/bmaupin/epub-samples/blob/master/minimal-v32/mimetype
// Spec: http://www.idpf.org/epub/301/spec/epub-ocf.html#sec-zip-container-mime
func writeMimetype(tempDir string) error {
	mimetypeFilePath := filepath.Join(tempDir, mimetypeFilename)

	if err := ioutil.WriteFile(mimetypeFilePath, []byte(mediaTypeEpub), filePermissions); err != nil {
		return err
	}

	return nil
}

func (e *Epub) writePackageFile(tempDir string) error {
	err := e.pkg.write(tempDir)
	if err != nil {
		return err
	}

	return nil
}

// Write the section files to the temporary directory and add the sections to
// the TOC and package files
func (e *Epub) writeSections(tempDir string) error {
	for i, section := range e.sections {
		sectionIndex := i + 1
		sectionFilename := fmt.Sprintf(sectionFileFormat, sectionIndex)
		sectionFilePath := filepath.Join(tempDir, contentFolderName, xhtmlFolderName, sectionFilename)
		section.write(sectionFilePath)

		relativePath := filepath.Join(xhtmlFolderName, sectionFilename)
		e.toc.addSection(sectionIndex, section.Title(), relativePath)
		e.pkg.addToManifest(sectionFilename, relativePath, mediaTypeXhtml, "")
		e.pkg.addToSpine(sectionFilename)
	}

	return nil
}

// Write the TOC file to the temporary directory and add the TOC entries to the
// package file
func (e *Epub) writeToc(tempDir string) error {
	e.pkg.addToManifest(tocNavItemID, tocNavFilename, mediaTypeXhtml, tocNavItemProperties)
	e.pkg.addToManifest(tocNcxItemID, tocNcxFilename, mediaTypeNcx, "")

	err := e.toc.write(tempDir)
	if err != nil {
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>package db

import (
	"bytes"
	"time"

	"github.com/boltdb/bolt"

	"github.com/subutai-io/base/agent/log"
)

var (
	bucket = "MyBucket"
	search = "SearchIndex"
	users  = "Users"
	tokens = "Tokens"
	authid = "AuthID"
	db     = initdb()
)

func initdb() *bolt.DB {
	db, err := bolt.Open("my.db", 0600, nil)
	log.Check(log.FatalLevel, "Openning DB: my.db", err)
	err = db.Update(func(tx *bolt.Tx) error {
		_, err := tx.CreateBucketIfNotExists([]byte(bucket))
		log.Check(log.FatalLevel, "Creating data bucket: "+bucket, err)

		_, err = tx.CreateBucketIfNotExists([]byte(search))
		log.Check(log.FatalLevel, "Creating search bucket: "+search, err)

		_, err = tx.CreateBucketIfNotExists([]byte(users))
		log.Check(log.FatalLevel, "Creating users bucket: "+search, err)

		_, err = tx.CreateBucketIfNotExists([]byte(tokens))
		log.Check(log.FatalLevel, "Creating users bucket: "+search, err)

		_, err = tx.CreateBucketIfNotExists([]byte(authid))
		log.Check(log.FatalLevel, "Creating users bucket: "+search, err)

		return nil
	})
	log.Check(log.FatalLevel, "Finishing update transaction"+bucket, err)

	return db
}

func Write(owner, key, value string, options ...map[string]string) {
	if len(owner) == 0 {
		owner = "subutai"
	}
	now, _ := time.Now().MarshalText()
	err := db.Update(func(tx *bolt.Tx) error {

		b, err := tx.Bucket([]byte(users)).CreateBucketIfNotExists([]byte(owner))
		log.Check(log.FatalLevel, "Creating users subbucket: "+key, err)
		b, err = b.CreateBucketIfNotExists([]byte("files"))
		log.Check(log.FatalLevel, "Creating users:files subbucket: "+key, err)
		b.Put([]byte(key), []byte(value))

		b, err = tx.Bucket([]byte(bucket)).CreateBucketIfNotExists([]byte(key))
		log.Check(log.FatalLevel, "Creating subbucket: "+key, err)
		b.Put([]byte("date"), now)
		b.Put([]byte("name"), []byte(value))
		b.Put([]byte("owner"), []byte(owner))

		for i, _ := range options {
			for k, v := range options[i] {
				err = b.Put([]byte(k), []byte(v))
				log.Check(log.WarnLevel, "Storing key: "+k, err)
			}
		}

		b, err = tx.Bucket([]byte(search)).CreateBucketIfNotExists([]byte(value))
		log.Check(log.FatalLevel, "Creating subbucket: "+key, err)
		b.Put(now, []byte(key))

		return nil
	})
	if err != nil {
		log.Fatal(err.Error())
	}
}

func Read(key string) (value string) {
	db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(bucket)).Bucket([]byte(key))
		value = string(b.Get([]byte("name")))
		return nil
	})
	return value
}

func List() map[string]string {
	list := make(map[string]string)
	db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(bucket))
		b.ForEach(func(k, v []byte) error {
			list[string(k)] = string(b.Bucket(k).Get([]byte("name")))
			return nil
		})
		return nil
	})
	return list
}

func Close() {
	db.Close()
}

func Search(query string) map[string]string {
	list := make(map[string]string)
	db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(search))
		c := b.Cursor()
		for k, _ := c.Seek([]byte(query)); bytes.HasPrefix(k, []byte(query)); k, _ = c.Next() {
			b.Bucket(k).ForEach(func(kk, vv []byte) error {
				list[string(vv)] = string(k)
				return nil
			})
		}
		return nil
	})
	return list
}

func LastHash(name string) (hash string) {
	db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(search)).Bucket([]byte(name))
		_, v := b.Cursor().Last()
		hash = string(v)
		return nil
	})
	return hash
}

func RegisterUser(name, key []byte) {
	db.Update(func(tx *bolt.Tx) error {

		b, err := tx.Bucket([]byte(users)).CreateBucketIfNotExists([]byte(name))
		log.Check(log.FatalLevel, "Creating users subbucket: "+string(name), err)
		b.Put([]byte("key"), key)

		return nil
	})
}

func UserKey(name string) (key string) {
	db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(users)).Bucket([]byte(name))
		key = string(b.Get([]byte("key")))
		return nil
	})
	return key
}

func SaveToken(name, token string) {
	db.Update(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(tokens))
		b.Put([]byte(token), []byte(name))
		return nil
	})
}

func CheckToken(token string) (name string) {
	db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(tokens))
		name = string(b.Get([]byte(token)))
		return nil
	})
	return name
}

func SaveAuthID(name, token string) {
	db.Update(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(authid))
		b.Put([]byte(token), []byte(name))
		return nil
	})
}

func CheckAuthID(token string) (name string) {
	db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(authid))
		name = string(b.Get([]byte(token)))
		return nil
	})
	return name
}
<commit_msg>Fixed panic on validate token.<commit_after>package db

import (
	"bytes"
	"time"

	"github.com/boltdb/bolt"

	"github.com/subutai-io/base/agent/log"
)

var (
	bucket = "MyBucket"
	search = "SearchIndex"
	users  = "Users"
	tokens = "Tokens"
	authid = "AuthID"
	db     = initdb()
)

func initdb() *bolt.DB {
	db, err := bolt.Open("my.db", 0600, nil)
	log.Check(log.FatalLevel, "Openning DB: my.db", err)
	err = db.Update(func(tx *bolt.Tx) error {
		_, err := tx.CreateBucketIfNotExists([]byte(bucket))
		log.Check(log.FatalLevel, "Creating data bucket: "+bucket, err)

		_, err = tx.CreateBucketIfNotExists([]byte(search))
		log.Check(log.FatalLevel, "Creating search bucket: "+search, err)

		_, err = tx.CreateBucketIfNotExists([]byte(users))
		log.Check(log.FatalLevel, "Creating users bucket: "+search, err)

		_, err = tx.CreateBucketIfNotExists([]byte(tokens))
		log.Check(log.FatalLevel, "Creating users bucket: "+search, err)

		_, err = tx.CreateBucketIfNotExists([]byte(authid))
		log.Check(log.FatalLevel, "Creating users bucket: "+search, err)

		return nil
	})
	log.Check(log.FatalLevel, "Finishing update transaction"+bucket, err)

	return db
}

func Write(owner, key, value string, options ...map[string]string) {
	if len(owner) == 0 {
		owner = "subutai"
	}
	now, _ := time.Now().MarshalText()
	err := db.Update(func(tx *bolt.Tx) error {

		b, err := tx.Bucket([]byte(users)).CreateBucketIfNotExists([]byte(owner))
		log.Check(log.FatalLevel, "Creating users subbucket: "+key, err)
		b, err = b.CreateBucketIfNotExists([]byte("files"))
		log.Check(log.FatalLevel, "Creating users:files subbucket: "+key, err)
		b.Put([]byte(key), []byte(value))

		b, err = tx.Bucket([]byte(bucket)).CreateBucketIfNotExists([]byte(key))
		log.Check(log.FatalLevel, "Creating subbucket: "+key, err)
		b.Put([]byte("date"), now)
		b.Put([]byte("name"), []byte(value))
		b.Put([]byte("owner"), []byte(owner))

		for i, _ := range options {
			for k, v := range options[i] {
				err = b.Put([]byte(k), []byte(v))
				log.Check(log.WarnLevel, "Storing key: "+k, err)
			}
		}

		b, err = tx.Bucket([]byte(search)).CreateBucketIfNotExists([]byte(value))
		log.Check(log.FatalLevel, "Creating subbucket: "+key, err)
		b.Put(now, []byte(key))

		return nil
	})
	if err != nil {
		log.Fatal(err.Error())
	}
}

func Read(key string) (value string) {
	db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(bucket)).Bucket([]byte(key))
		value = string(b.Get([]byte("name")))
		return nil
	})
	return value
}

func List() map[string]string {
	list := make(map[string]string)
	db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(bucket))
		b.ForEach(func(k, v []byte) error {
			list[string(k)] = string(b.Bucket(k).Get([]byte("name")))
			return nil
		})
		return nil
	})
	return list
}

func Close() {
	db.Close()
}

func Search(query string) map[string]string {
	list := make(map[string]string)
	db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(search))
		c := b.Cursor()
		for k, _ := c.Seek([]byte(query)); bytes.HasPrefix(k, []byte(query)); k, _ = c.Next() {
			b.Bucket(k).ForEach(func(kk, vv []byte) error {
				list[string(vv)] = string(k)
				return nil
			})
		}
		return nil
	})
	return list
}

func LastHash(name string) (hash string) {
	db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(search)).Bucket([]byte(name))
		_, v := b.Cursor().Last()
		hash = string(v)
		return nil
	})
	return hash
}

func RegisterUser(name, key []byte) {
	db.Update(func(tx *bolt.Tx) error {

		b, err := tx.Bucket([]byte(users)).CreateBucketIfNotExists([]byte(name))
		log.Check(log.FatalLevel, "Creating users subbucket: "+string(name), err)
		b.Put([]byte("key"), key)

		return nil
	})
}

func UserKey(name string) (key string) {
	db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(users)).Bucket([]byte(name))
		key = string(b.Get([]byte("key")))
		return nil
	})
	return key
}

func SaveToken(name, token string) {
	db.Update(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(tokens))
		b.Put([]byte(token), []byte(name))
		return nil
	})
}

func CheckToken(token string) (name string) {
	db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(tokens))
		if value := b.Get([]byte(token)); value != nil {
			name = string(value)
		}
		return nil
	})
	return name
}

func SaveAuthID(name, token string) {
	db.Update(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(authid))
		b.Put([]byte(token), []byte(name))
		return nil
	})
}

func CheckAuthID(token string) (name string) {
	db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(authid))
		name = string(b.Get([]byte(token)))
		return nil
	})
	return name
}
<|endoftext|>


<commit_before>package db

import "fmt"

//----------------------------------------
// Main entry

type dbBackendType string

const (
	LevelDBBackend   dbBackendType = "leveldb" // legacy, defaults to goleveldb unless +gcc
	CLevelDBBackend  dbBackendType = "cleveldb"
	GoLevelDBBackend dbBackendType = "goleveldb"
	MemDBBackend     dbBackendType = "memDB"
	FSDBBackend      dbBackendType = "fsdb" // using the filesystem naively
)

type dbCreator func(name string, dir string) (DB, error)

var backends = map[dbBackendType]dbCreator{}

func registerDBCreator(backend dbBackendType, creator dbCreator, force bool) {
	_, ok := backends[backend]
	if !force && ok {
		return
	}
	backends[backend] = creator
}

func NewDB(name string, backend dbBackendType, dir string) DB {
	db, err := backends[backend](name, dir)
	if err != nil {
		panic(fmt.Sprintf("Error initializing DB: %v", err))
	}
	return db
}
<commit_msg>lowercase memDB type key<commit_after>package db

import "fmt"

//----------------------------------------
// Main entry

type dbBackendType string

const (
	LevelDBBackend   dbBackendType = "leveldb" // legacy, defaults to goleveldb unless +gcc
	CLevelDBBackend  dbBackendType = "cleveldb"
	GoLevelDBBackend dbBackendType = "goleveldb"
	MemDBBackend     dbBackendType = "memdb"
	FSDBBackend      dbBackendType = "fsdb" // using the filesystem naively
)

type dbCreator func(name string, dir string) (DB, error)

var backends = map[dbBackendType]dbCreator{}

func registerDBCreator(backend dbBackendType, creator dbCreator, force bool) {
	_, ok := backends[backend]
	if !force && ok {
		return
	}
	backends[backend] = creator
}

func NewDB(name string, backend dbBackendType, dir string) DB {
	db, err := backends[backend](name, dir)
	if err != nil {
		panic(fmt.Sprintf("Error initializing DB: %v", err))
	}
	return db
}
<|endoftext|>


<commit_before>// go run debug.go [URL]

// +build ignore

package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"

	"github.com/SlyMarbo/rss"
)

func main() {
	if len(os.Args) == 1 {
		fmt.Fprintf(os.Stderr, "Usage:\n\t%s [URL]\n", filepath.Base(os.Args[0]))
		os.Exit(2)
	}

	feed, err := rss.Fetch(os.Args[1])
	if err != nil {
		panic(err)
	}

	raw, err := json.Marshal(feed)
	if err != nil {
		panic(err)
	}

	buf := new(bytes.Buffer)
	if err := json.Indent(buf, raw, "", "\t"); err != nil {
		panic(err)
	}

	fmt.Println(buf.String())
}
<commit_msg>Improved usage message for debug script<commit_after>// go run debug.go [URL]

// +build ignore

package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"

	"github.com/SlyMarbo/rss"
)

func main() {
	if len(os.Args) != 2 || os.Args[1] == "-h" {
		fmt.Fprintf(os.Stderr, "Usage:\n\t%s [URL]\n", filepath.Base(os.Args[0]))
		os.Exit(2)
	}

	feed, err := rss.Fetch(os.Args[1])
	if err != nil {
		panic(err)
	}

	raw, err := json.Marshal(feed)
	if err != nil {
		panic(err)
	}

	buf := new(bytes.Buffer)
	if err := json.Indent(buf, raw, "", "\t"); err != nil {
		panic(err)
	}

	fmt.Println(buf.String())
}
<|endoftext|>