<commit_before>package security

import (
	"fmt"
	"time"

	"github.com/lestrrat-go/jwx/jwk"
	"github.com/sirupsen/logrus"
)

var jwkCache map[string]*jwk.Set

func init() {
	jwkCache = map[string]*jwk.Set{}
}

func getJWKByEndpoint(endpoint, keyID string) (interface{}, error) {
	keys := jwkCache[endpoint]
	if keys == nil {
		return nil, fmt.Errorf("No keys found for endpoint %s", endpoint)
	}
	k := keys.LookupKeyID(keyID)
	if len(k) == 0 {
		loadEndpoint(endpoint)
		keys = jwkCache[endpoint]
		k = keys.LookupKeyID(keyID)
		if len(k) == 0 {
			return nil, fmt.Errorf("Key %s not found at %s", keyID, endpoint)
		}
	}
	if len(k) > 1 {
		return nil, fmt.Errorf("Unexpected error, more than 1 key %s found at %s", keyID, endpoint)
	}
	return k[0].Materialize()
}

// LoadAllEndpoints loads all the endpoints
func LoadAllEndpoints(endpoints []string) {
	for _, ep := range endpoints {
		loadEndpoint(ep)
	}
}

// RefreshEndpoints refreshes endpoints
func RefreshEndpoints(endpoints []string) {
	ticker := time.NewTicker(time.Hour * 24)
	go func() {
		for true {
			<-ticker.C
			for _, ep := range endpoints {
				loadEndpoint(ep)
			}
		}
	}()
}

func loadEndpoint(endpoint string) {
	keySet, err := jwk.FetchHTTP(endpoint)
	if err != nil {
		logrus.WithError(err).WithField("endpoint", endpoint).Error("Unable to load keys for endpoint")
	}
	jwkCache[endpoint] = keySet
}
<commit_msg>[tweek-gateway] reload jwk keys on error (#1193)<commit_after>package security

import (
	"fmt"
	"time"

	"github.com/lestrrat-go/jwx/jwk"
	"github.com/sirupsen/logrus"
)

type jwkRecord struct {
	set *jwk.Set
	err error
}

var jwkCache map[string]*jwkRecord

func init() {
	jwkCache = map[string]*jwkRecord{}
}

func getJWKByEndpoint(endpoint, keyID string) (interface{}, error) {
	rec, ok := jwkCache[endpoint]
	if !ok {
		return nil, fmt.Errorf("No keys found for endpoint %s", endpoint)
	}

	if rec.err != nil {
		return nil, rec.err
	}

	k := rec.set.LookupKeyID(keyID)
	if len(k) == 0 {
		rec := loadEndpoint(endpoint)
		if rec.err != nil {
			return nil, rec.err
		}
		k = rec.set.LookupKeyID(keyID)
		if len(k) == 0 {
			return nil, fmt.Errorf("Key %s not found at %s", keyID, endpoint)
		}
	}
	if len(k) > 1 {
		return nil, fmt.Errorf("Unexpected error, more than 1 key %s found at %s", keyID, endpoint)
	}
	return k[0].Materialize()
}

// LoadAllEndpoints loads all the endpoints
func LoadAllEndpoints(endpoints []string) {
	for _, ep := range endpoints {
		loadEndpoint(ep)
	}
}

// RefreshEndpoints refreshes endpoints
func RefreshEndpoints(endpoints []string) {
	ticker := time.NewTicker(time.Hour * 24)
	go func() {
		for true {
			<-ticker.C
			for _, ep := range endpoints {
				loadEndpoint(ep)
			}
		}
	}()
}

func loadEndpoint(endpoint string) *jwkRecord {
	return loadEndpointWithRetry(endpoint, 0)
}

func loadEndpointWithRetry(endpoint string, retryCount uint) *jwkRecord {
	rec := &jwkRecord{}
	rec.set, rec.err = jwk.FetchHTTP(endpoint)
	jwkCache[endpoint] = rec

	if rec.err != nil {
		logrus.WithError(rec.err).WithField("endpoint", endpoint).Error("Unable to load keys for endpoint")

		go func() {
			<-time.After(time.Second * (1 << retryCount))
			cached := jwkCache[endpoint]
			if cached == rec {
				nextCount := retryCount + 1
				if nextCount > 6 {
					// limit retry delay to 64 Seconds (~1 Minute)
					nextCount = 6
				}
				loadEndpointWithRetry(endpoint, nextCount)
			}
		}()
	}

	return rec
}
<|endoftext|>


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

package testlapack

import (
	"fmt"
	"math"
	"testing"

	"golang.org/x/exp/rand"
)

type Dlanv2er interface {
	Dlanv2(a, b, c, d float64) (aa, bb, cc, dd float64, rt1r, rt1i, rt2r, rt2i float64, cs, sn float64)
}

func Dlanv2Test(t *testing.T, impl Dlanv2er) {
	rnd := rand.New(rand.NewSource(1))
	for i := 0; i < 1000; i++ {
		a := rnd.NormFloat64()
		b := rnd.NormFloat64()
		c := rnd.NormFloat64()
		d := rnd.NormFloat64()
		aa, bb, cc, dd, rt1r, rt1i, rt2r, rt2i, cs, sn := impl.Dlanv2(a, b, c, d)

		mat := fmt.Sprintf("[%v %v; %v %v]", a, b, c, d)
		if cc == 0 {
			if rt1i != 0 || rt2i != 0 {
				t.Errorf("Unexpected complex eigenvalues for %v", mat)
			}
		} else {
			if aa != dd {
				t.Errorf("Diagonal elements not equal for %v: got [%v %v]", mat, aa, dd)
			}
			if bb*cc >= 0 {
				t.Errorf("Non-diagonal elements have the same sign for %v: got [%v %v]", mat, bb, cc)
			} else {
				im := math.Sqrt(-bb * cc)
				if math.Abs(rt1i-im) > 1e-14 && math.Abs(rt1i+im) > 1e-14 {
					t.Errorf("Unexpected imaginary part of eigenvalue for %v: got %v, want %v or %v", mat, rt1i, im, -im)
				}
				if math.Abs(rt2i-im) > 1e-14 && math.Abs(rt2i+im) > 1e-14 {
					t.Errorf("Unexpected imaginary part of eigenvalue for %v: got %v, want %v or %v", mat, rt2i, im, -im)
				}
			}
		}
		if rt1r != aa && rt1r != dd {
			t.Errorf("Unexpected real part of eigenvalue for %v: got %v, want %v or %v", mat, rt1r, aa, dd)
		}
		if rt2r != aa && rt2r != dd {
			t.Errorf("Unexpected real part of eigenvalue for %v: got %v, want %v or %v", mat, rt2r, aa, dd)
		}
		if math.Abs(math.Hypot(cs, sn)-1) > 1e-14 {
			t.Errorf("Unexpected unitary matrix for %v: got cs %v, sn %v", mat, cs, sn)
		}

		gota := cs*(aa*cs-bb*sn) - sn*(cc*cs-dd*sn)
		gotb := cs*(aa*sn+bb*cs) - sn*(cc*sn+dd*cs)
		gotc := sn*(aa*cs-bb*sn) + cs*(cc*cs-dd*sn)
		gotd := sn*(aa*sn+bb*cs) + cs*(cc*sn+dd*cs)
		if math.Abs(gota-a) > 1e-14 ||
			math.Abs(gotb-b) > 1e-14 ||
			math.Abs(gotc-c) > 1e-14 ||
			math.Abs(gotd-d) > 1e-14 {
			t.Errorf("Unexpected factorization: got [%v %v; %v %v], want [%v %v; %v %v]", gota, gotb, gotc, gotd, a, b, c, d)
		}
	}
}
<commit_msg>lapack/testlapack: reduce indentation in Dlanv2 test<commit_after>// Copyright ©2016 The Gonum Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package testlapack

import (
	"fmt"
	"math"
	"testing"

	"golang.org/x/exp/rand"
)

type Dlanv2er interface {
	Dlanv2(a, b, c, d float64) (aa, bb, cc, dd float64, rt1r, rt1i, rt2r, rt2i float64, cs, sn float64)
}

func Dlanv2Test(t *testing.T, impl Dlanv2er) {
	rnd := rand.New(rand.NewSource(1))
	for i := 0; i < 1000; i++ {
		a := rnd.NormFloat64()
		b := rnd.NormFloat64()
		c := rnd.NormFloat64()
		d := rnd.NormFloat64()
		dlanv2Test(t, impl, a, b, c, d)
	}
}

func dlanv2Test(t *testing.T, impl Dlanv2er, a, b, c, d float64) {
	aa, bb, cc, dd, rt1r, rt1i, rt2r, rt2i, cs, sn := impl.Dlanv2(a, b, c, d)

	mat := fmt.Sprintf("[%v %v; %v %v]", a, b, c, d)
	if cc == 0 {
		if rt1i != 0 || rt2i != 0 {
			t.Errorf("Unexpected complex eigenvalues for %v", mat)
		}
	} else {
		if aa != dd {
			t.Errorf("Diagonal elements not equal for %v: got [%v %v]", mat, aa, dd)
		}
		if bb*cc >= 0 {
			t.Errorf("Non-diagonal elements have the same sign for %v: got [%v %v]", mat, bb, cc)
		} else {
			im := math.Sqrt(-bb * cc)
			if math.Abs(rt1i-im) > 1e-14 && math.Abs(rt1i+im) > 1e-14 {
				t.Errorf("Unexpected imaginary part of eigenvalue for %v: got %v, want %v or %v", mat, rt1i, im, -im)
			}
			if math.Abs(rt2i-im) > 1e-14 && math.Abs(rt2i+im) > 1e-14 {
				t.Errorf("Unexpected imaginary part of eigenvalue for %v: got %v, want %v or %v", mat, rt2i, im, -im)
			}
		}
	}
	if rt1r != aa && rt1r != dd {
		t.Errorf("Unexpected real part of eigenvalue for %v: got %v, want %v or %v", mat, rt1r, aa, dd)
	}
	if rt2r != aa && rt2r != dd {
		t.Errorf("Unexpected real part of eigenvalue for %v: got %v, want %v or %v", mat, rt2r, aa, dd)
	}
	if math.Abs(math.Hypot(cs, sn)-1) > 1e-14 {
		t.Errorf("Unexpected unitary matrix for %v: got cs %v, sn %v", mat, cs, sn)
	}

	gota := cs*(aa*cs-bb*sn) - sn*(cc*cs-dd*sn)
	gotb := cs*(aa*sn+bb*cs) - sn*(cc*sn+dd*cs)
	gotc := sn*(aa*cs-bb*sn) + cs*(cc*cs-dd*sn)
	gotd := sn*(aa*sn+bb*cs) + cs*(cc*sn+dd*cs)
	if math.Abs(gota-a) > 1e-14 ||
		math.Abs(gotb-b) > 1e-14 ||
		math.Abs(gotc-c) > 1e-14 ||
		math.Abs(gotd-d) > 1e-14 {
		t.Errorf("Unexpected factorization: got [%v %v; %v %v], want [%v %v; %v %v]", gota, gotb, gotc, gotd, a, b, c, d)
	}
}
<|endoftext|>


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

package result

import (
	"bytes"
	"gopkg.in/goyy/goyy.v0/data/domain"
	"gopkg.in/goyy/goyy.v0/data/entity"
	"gopkg.in/goyy/goyy.v0/util/errors"
	"gopkg.in/goyy/goyy.v0/util/jsons"
	"gopkg.in/goyy/goyy.v0/util/strings"
	"strconv"
)

type Page struct {
	Success bool        `json:"success"`
	Id      string      `json:"id"`
	Token   string      `json:"token"`
	Code    string      `json:"code"`
	Message string      `json:"message"`
	Memo    string      `json:"memo"`
	Tag     string      `json:"tag"`
	Data    domain.Page `json:"data"`
}

func (me *Page) JSON() string {
	var b bytes.Buffer
	b.WriteString(`{"success":` + strconv.FormatBool(me.Success) + ",")
	b.WriteString(`"id":"` + jsons.Format(me.Id) + `",`)
	b.WriteString(`"token":"` + jsons.Format(me.Token) + `",`)
	b.WriteString(`"code":"` + jsons.Format(me.Code) + `",`)
	b.WriteString(`"message":"` + jsons.Format(me.Message) + `",`)
	b.WriteString(`"memo":"` + jsons.Format(me.Memo) + `",`)
	b.WriteString(`"tag":"` + jsons.Format(me.Tag) + `",`)
	b.WriteString(`"data":{`)
	if me.Data != nil {
		b.WriteString(`"pageNo":` + strconv.Itoa(me.Data.PageNo()) + `,`)
		b.WriteString(`"pageSize":` + strconv.Itoa(me.Data.PageSize()) + `,`)
		b.WriteString(`"totalElements":` + strconv.Itoa(me.Data.TotalElements()) + `,`)
		b.WriteString(`"function":"` + me.Data.Function() + `",`)
		b.WriteString(`"length":` + strconv.Itoa(me.Data.Length()) + `,`)
		b.WriteString(`"slider":` + strconv.Itoa(me.Data.Slider()) + `,`)
		b.WriteString(`"slice":[`)
		for i := 0; i < me.Data.Content().Len(); i++ {
			if i > 0 {
				b.WriteString(",")
			}
			b.WriteString(entity.FormatJSON(me.Data.Content().Index(i)))
		}
		b.WriteString("]")
	}
	b.WriteString("}}")
	return b.String()
}

func (me *Page) ParseJSON(json string) error {
	if strings.IsBlank(json) {
		return nil
	}
	// JSON format validation
	if !strings.HasPrefix(json, `{"success":`) || !strings.HasSuffix(json, "}]}}") {
		return errors.New("JSON format is not legal : no success")
	}
	if !strings.Contains(json, `,"data":{`) {
		return errors.New("JSON format is not legal : no data")
	}
	if !strings.Contains(json, `"pageNo":`) {
		return errors.New("JSON format is not legal : no pageNo")
	}
	if !strings.Contains(json, `"pageSize":`) {
		return errors.New("JSON format is not legal : no pageSize")
	}
	if !strings.Contains(json, `"totalElements":`) {
		return errors.New("JSON format is not legal : no totalElements")
	}
	if !strings.Contains(json, `,"slice":[{`) {
		return errors.New("JSON format is not legal : no slice")
	}
	// result info
	if "true" == strings.Between(json, `"success":`, `,"`) {
		me.Success = true
	}
	content := jreplace(json)
	me.Code = jparse(strings.Between(content, `"code":"`, `",`))
	me.Message = jparse(strings.Between(content, `"message":"`, `",`))
	me.Memo = jparse(strings.Between(content, `"memo":"`, `",`))
	me.Tag = jparse(strings.Between(content, `"tag":"`, `",`))
	// page info
	pagejson := strings.After(json, `,"data":{`)
	pagejson = strings.BeforeLast(pagejson, "}}")
	strPageNo := strings.Between(pagejson, `"pageNo":`, `,"`)
	strPageSize := strings.Between(pagejson, `"pageSize":`, `,"`)
	strTotalElements := strings.Between(pagejson, `"totalElements":`, `,"`)
	if pageNo, err := strconv.Atoi(strPageNo); err != nil {
		if pageSize, err := strconv.Atoi(strPageSize); err != nil {
			pageable := domain.NewPageable(pageNo, pageSize)
			me.Data.SetPageable(pageable)
		}
	}
	if totalElements, err := strconv.Atoi(strTotalElements); err != nil {
		me.Data.SetTotalElements(totalElements)
	}
	// slice info
	datajson := strings.Between(json, `,"slice":[{`, "}]}}")
	datas := strings.Split(datajson, "},{")
	for _, data := range datas {
		e := me.Data.Content().New()
		if err := entity.ParseJSON(e, "{"+data+"}"); err != nil {
			return err
		}
		me.Data.Content().Append(e)
	}
	return nil
}
<commit_msg>Page.JSON() : Add totalPages<commit_after>// Copyright 2014 The goyy Authors.  All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

package result

import (
	"bytes"
	"gopkg.in/goyy/goyy.v0/data/domain"
	"gopkg.in/goyy/goyy.v0/data/entity"
	"gopkg.in/goyy/goyy.v0/util/errors"
	"gopkg.in/goyy/goyy.v0/util/jsons"
	"gopkg.in/goyy/goyy.v0/util/strings"
	"strconv"
)

type Page struct {
	Success bool        `json:"success"`
	Id      string      `json:"id"`
	Token   string      `json:"token"`
	Code    string      `json:"code"`
	Message string      `json:"message"`
	Memo    string      `json:"memo"`
	Tag     string      `json:"tag"`
	Data    domain.Page `json:"data"`
}

func (me *Page) JSON() string {
	var b bytes.Buffer
	b.WriteString(`{"success":` + strconv.FormatBool(me.Success) + ",")
	b.WriteString(`"id":"` + jsons.Format(me.Id) + `",`)
	b.WriteString(`"token":"` + jsons.Format(me.Token) + `",`)
	b.WriteString(`"code":"` + jsons.Format(me.Code) + `",`)
	b.WriteString(`"message":"` + jsons.Format(me.Message) + `",`)
	b.WriteString(`"memo":"` + jsons.Format(me.Memo) + `",`)
	b.WriteString(`"tag":"` + jsons.Format(me.Tag) + `",`)
	b.WriteString(`"data":{`)
	if me.Data != nil {
		b.WriteString(`"pageNo":` + strconv.Itoa(me.Data.PageNo()) + `,`)
		b.WriteString(`"pageSize":` + strconv.Itoa(me.Data.PageSize()) + `,`)
		b.WriteString(`"totalPages":` + strconv.Itoa(me.Data.TotalPages()) + `,`)
		b.WriteString(`"totalElements":` + strconv.Itoa(me.Data.TotalElements()) + `,`)
		b.WriteString(`"function":"` + me.Data.Function() + `",`)
		b.WriteString(`"length":` + strconv.Itoa(me.Data.Length()) + `,`)
		b.WriteString(`"slider":` + strconv.Itoa(me.Data.Slider()) + `,`)
		b.WriteString(`"slice":[`)
		for i := 0; i < me.Data.Content().Len(); i++ {
			if i > 0 {
				b.WriteString(",")
			}
			b.WriteString(entity.FormatJSON(me.Data.Content().Index(i)))
		}
		b.WriteString("]")
	}
	b.WriteString("}}")
	return b.String()
}

func (me *Page) ParseJSON(json string) error {
	if strings.IsBlank(json) {
		return nil
	}
	// JSON format validation
	if !strings.HasPrefix(json, `{"success":`) || !strings.HasSuffix(json, "}]}}") {
		return errors.New("JSON format is not legal : no success")
	}
	if !strings.Contains(json, `,"data":{`) {
		return errors.New("JSON format is not legal : no data")
	}
	if !strings.Contains(json, `"pageNo":`) {
		return errors.New("JSON format is not legal : no pageNo")
	}
	if !strings.Contains(json, `"pageSize":`) {
		return errors.New("JSON format is not legal : no pageSize")
	}
	if !strings.Contains(json, `"totalElements":`) {
		return errors.New("JSON format is not legal : no totalElements")
	}
	if !strings.Contains(json, `,"slice":[{`) {
		return errors.New("JSON format is not legal : no slice")
	}
	// result info
	if "true" == strings.Between(json, `"success":`, `,"`) {
		me.Success = true
	}
	content := jreplace(json)
	me.Code = jparse(strings.Between(content, `"code":"`, `",`))
	me.Message = jparse(strings.Between(content, `"message":"`, `",`))
	me.Memo = jparse(strings.Between(content, `"memo":"`, `",`))
	me.Tag = jparse(strings.Between(content, `"tag":"`, `",`))
	// page info
	pagejson := strings.After(json, `,"data":{`)
	pagejson = strings.BeforeLast(pagejson, "}}")
	strPageNo := strings.Between(pagejson, `"pageNo":`, `,"`)
	strPageSize := strings.Between(pagejson, `"pageSize":`, `,"`)
	strTotalElements := strings.Between(pagejson, `"totalElements":`, `,"`)
	if pageNo, err := strconv.Atoi(strPageNo); err != nil {
		if pageSize, err := strconv.Atoi(strPageSize); err != nil {
			pageable := domain.NewPageable(pageNo, pageSize)
			me.Data.SetPageable(pageable)
		}
	}
	if totalElements, err := strconv.Atoi(strTotalElements); err != nil {
		me.Data.SetTotalElements(totalElements)
	}
	// slice info
	datajson := strings.Between(json, `,"slice":[{`, "}]}}")
	datas := strings.Split(datajson, "},{")
	for _, data := range datas {
		e := me.Data.Content().New()
		if err := entity.ParseJSON(e, "{"+data+"}"); err != nil {
			return err
		}
		me.Data.Content().Append(e)
	}
	return nil
}
<|endoftext|>


<commit_before>package gengorums

import (
	"bytes"
	"fmt"
	"go/ast"
	"go/build"
	"go/format"
	"go/parser"
	"go/printer"
	"go/token"
	"go/types"
	"io/ioutil"
	"log"
	"os"
	"path"
	"sort"
	"strings"
	"unicode"

	"github.com/google/go-cmp/cmp"
	"golang.org/x/tools/go/loader"
)

const (
	devPkgPath = "./cmd/protoc-gen-gorums/dev"
)

// GenerateBundleFile generates a file with static definitions for Gorums.
func GenerateBundleFile(dst string) {
	pkgIdent, reservedIdents, code := bundle(devPkgPath)
	// escape backticks
	escaped := strings.ReplaceAll(string(code), "`", "`+\"`\"+`")
	src := "// Code generated by protoc-gen-gorums. DO NOT EDIT.\n" +
		"// Source files can be found in: " + devPkgPath + "\n\n" +
		"package gengorums\n\n" +
		generatePkgMap(pkgIdent, reservedIdents) +
		"var staticCode = `" + escaped + "`\n"

	staticContent, err := format.Source([]byte(src))
	if err != nil {
		log.Fatalf("formatting failed: %v", err)
	}
	currentContent, err := ioutil.ReadFile(dst)
	if err != nil {
		log.Fatal(err)
	}
	if diff := cmp.Diff(currentContent, staticContent); diff != "" {
		fmt.Fprintf(os.Stderr, "change detected (-current +new):\n%s", diff)
		fmt.Fprintf(os.Stderr, "\nReview changes above; to revert use:\n")
		fmt.Fprintf(os.Stderr, "mv %s.bak %s\n", dst, dst)
	}
	err = ioutil.WriteFile(dst, []byte(staticContent), 0666)
	if err != nil {
		log.Fatal(err)
	}
}

func generatePkgMap(pkgs map[string]string, reservedIdents []string) string {
	var keys []string
	for k := range pkgs {
		keys = append(keys, k)
	}
	sort.Strings(keys)
	var buf bytes.Buffer
	buf.WriteString(`
	// pkgIdentMap maps from package name to one of the package's identifiers.
	// These identifiers are used by the Gorums protoc plugin to generate
	// appropriate import statements.
	`)
	buf.WriteString("var pkgIdentMap = map[string]string{\n")
	for _, imp := range keys {
		buf.WriteString("\t" + `"`)
		buf.WriteString(imp)
		buf.WriteString(`": "`)
		buf.WriteString(pkgs[imp])
		buf.WriteString(`",` + "\n")
	}
	buf.WriteString("}\n\n")
	buf.WriteString(`
	// reservedIdents holds the set of Gorums reserved identifiers.
	// These identifiers cannot be used to define message types in a proto file.
	`)
	buf.WriteString("var reservedIdents = []string{\n")
	for _, ident := range reservedIdents {
		buf.WriteString("\t" + `"`)
		buf.WriteString(ident)
		buf.WriteString(`",` + "\n")
	}
	buf.WriteString("}\n\n")
	return buf.String()
}

// findIdentifiers examines the given package to find all imported packages,
// and one used identifier in that imported package. These identifiers are
// used by the Gorums protoc plugin to generate appropriate import statements.
func findIdentifiers(fset *token.FileSet, info *loader.PackageInfo) (map[string]string, []string) {
	packagePath := info.Pkg.Path()
	pkgIdents := make(map[string][]string)
	for id, obj := range info.Uses {
		pos := fset.Position(id.Pos())
		if strings.Contains(pos.Filename, "zorums") {
			// ignore identifiers in zorums generated files
			continue
		}
		r := []rune(obj.Name())
		if unicode.IsLower(r[0]) {
			// ignore unexported identifiers
			continue
		}
		if pkg := obj.Pkg(); pkg != nil {
			switch obj := obj.(type) {
			case *types.Func:
				if typ := obj.Type(); typ != nil {
					if recv := typ.(*types.Signature).Recv(); recv != nil {
						// ignore functions on non-package types
						continue
					}
				}
				addUniqueIdentifier(pkgIdents, pkg.Path(), obj.Name())

			case *types.Const, *types.TypeName:
				addUniqueIdentifier(pkgIdents, pkg.Path(), obj.Name())

			case *types.Var:
				// no need to import obj's pkg if obj is a field of a struct.
				if obj.IsField() {
					continue
				}
				addUniqueIdentifier(pkgIdents, pkg.Path(), obj.Name())
			}
		}
	}
	// Only need to store one identifier for each imported package.
	// However, to ensure stable output, we sort the identifiers
	// and use the first element.
	pkgIdent := make(map[string]string, len(pkgIdents))
	var reservedIdents []string
	for path, idents := range pkgIdents {
		sort.Strings(idents)
		if path == packagePath {
			reservedIdents = idents
			continue
		}
		pkgIdent[path] = idents[0]
	}
	return pkgIdent, reservedIdents
}

func addUniqueIdentifier(pkgIdents map[string][]string, path, name string) {
	currentNames := pkgIdents[path]
	for _, known := range currentNames {
		if name == known {
			return
		}
	}
	pkgIdents[path] = append(pkgIdents[path], name)
}

// bundle returns a slice with the code for the given src package without imports.
// The returned map contains packages to be imported along with one identifier
// using the relevant import path.
// Loosly based on x/tools/cmd/bundle
func bundle(src string) (map[string]string, []string, []byte) {
	conf := loader.Config{ParserMode: parser.ParseComments, Build: &build.Default}
	conf.Import(src)
	lprog, err := conf.Load()
	if err != nil {
		log.Fatalf("failed to load Go package: %v", err)
	}
	// Because there was a single Import call and Load succeeded,
	// InitialPackages is guaranteed to hold the sole requested package.
	info := lprog.InitialPackages()[0]

	var out bytes.Buffer
	printFiles(&out, lprog.Fset, info)
	pkgIdentMap, reservedIdents := findIdentifiers(lprog.Fset, info)
	return pkgIdentMap, reservedIdents, out.Bytes()
}

func printFiles(out *bytes.Buffer, fset *token.FileSet, info *loader.PackageInfo) {
	for _, f := range info.Files {
		// filter files in dev package that shouldn't be bundled in template_static.go
		fileName := fset.File(f.Pos()).Name()
		if ignore(fileName) {
			continue
		}

		last := f.Package
		if len(f.Imports) > 0 {
			imp := f.Imports[len(f.Imports)-1]
			last = imp.End()
			if imp.Comment != nil {
				if e := imp.Comment.End(); e > last {
					last = e
				}
			}
		}

		// Pretty-print package-level declarations.
		// but no package or import declarations.
		var buf bytes.Buffer
		for _, decl := range f.Decls {
			if decl, ok := decl.(*ast.GenDecl); ok && decl.Tok == token.IMPORT {
				continue
			}
			beg, end := sourceRange(decl)
			printComments(out, f.Comments, last, beg)

			buf.Reset()
			format.Node(&buf, fset, &printer.CommentedNode{Node: decl, Comments: f.Comments})
			out.Write(buf.Bytes())
			last = printSameLineComment(out, f.Comments, fset, end)
			out.WriteString("\n\n")
		}
		printLastComments(out, f.Comments, last)
	}
}

// ignore files in dev folder with suffixes that shouldn't be bundled.
func ignore(file string) bool {
	for _, suffix := range []string{".proto", "_test.go"} {
		if strings.HasSuffix(file, suffix) {
			return true
		}
	}
	base := path.Base(file)
	if strings.HasPrefix(base, "zorums") {
		return true
	}
	return false
}

// sourceRange returns the [beg, end) interval of source code
// belonging to decl (incl. associated comments).
func sourceRange(decl ast.Decl) (beg, end token.Pos) {
	beg = decl.Pos()
	end = decl.End()

	var doc, com *ast.CommentGroup

	switch d := decl.(type) {
	case *ast.GenDecl:
		doc = d.Doc
		if len(d.Specs) > 0 {
			switch spec := d.Specs[len(d.Specs)-1].(type) {
			case *ast.ValueSpec:
				com = spec.Comment
			case *ast.TypeSpec:
				com = spec.Comment
			}
		}
	case *ast.FuncDecl:
		doc = d.Doc
	}

	if doc != nil {
		beg = doc.Pos()
	}
	if com != nil && com.End() > end {
		end = com.End()
	}

	return beg, end
}

func printPackageComments(out *bytes.Buffer, files []*ast.File) {
	// Concatenate package comments from all files...
	for _, f := range files {
		if doc := f.Doc.Text(); strings.TrimSpace(doc) != "" {
			for _, line := range strings.Split(doc, "\n") {
				fmt.Fprintf(out, "// %s\n", line)
			}
		}
	}
	// ...but don't let them become the actual package comment.
	fmt.Fprintln(out)
}

func printComments(out *bytes.Buffer, comments []*ast.CommentGroup, pos, end token.Pos) {
	for _, cg := range comments {
		if pos <= cg.Pos() && cg.Pos() < end {
			for _, c := range cg.List {
				fmt.Fprintln(out, c.Text)
			}
			fmt.Fprintln(out)
		}
	}
}

const infinity = 1 << 30

func printLastComments(out *bytes.Buffer, comments []*ast.CommentGroup, pos token.Pos) {
	printComments(out, comments, pos, infinity)
}

func printSameLineComment(out *bytes.Buffer, comments []*ast.CommentGroup, fset *token.FileSet, pos token.Pos) token.Pos {
	tf := fset.File(pos)
	for _, cg := range comments {
		if pos <= cg.Pos() && tf.Line(cg.Pos()) == tf.Line(pos) {
			for _, c := range cg.List {
				fmt.Fprintln(out, c.Text)
			}
			return cg.End()
		}
	}
	return pos
}
<commit_msg>simplified ignore check<commit_after>package gengorums

import (
	"bytes"
	"fmt"
	"go/ast"
	"go/build"
	"go/format"
	"go/parser"
	"go/printer"
	"go/token"
	"go/types"
	"io/ioutil"
	"log"
	"os"
	"path"
	"sort"
	"strings"
	"unicode"

	"github.com/google/go-cmp/cmp"
	"golang.org/x/tools/go/loader"
)

const (
	devPkgPath = "./cmd/protoc-gen-gorums/dev"
)

// GenerateBundleFile generates a file with static definitions for Gorums.
func GenerateBundleFile(dst string) {
	pkgIdent, reservedIdents, code := bundle(devPkgPath)
	// escape backticks
	escaped := strings.ReplaceAll(string(code), "`", "`+\"`\"+`")
	src := "// Code generated by protoc-gen-gorums. DO NOT EDIT.\n" +
		"// Source files can be found in: " + devPkgPath + "\n\n" +
		"package gengorums\n\n" +
		generatePkgMap(pkgIdent, reservedIdents) +
		"var staticCode = `" + escaped + "`\n"

	staticContent, err := format.Source([]byte(src))
	if err != nil {
		log.Fatalf("formatting failed: %v", err)
	}
	currentContent, err := ioutil.ReadFile(dst)
	if err != nil {
		log.Fatal(err)
	}
	if diff := cmp.Diff(currentContent, staticContent); diff != "" {
		fmt.Fprintf(os.Stderr, "change detected (-current +new):\n%s", diff)
		fmt.Fprintf(os.Stderr, "\nReview changes above; to revert use:\n")
		fmt.Fprintf(os.Stderr, "mv %s.bak %s\n", dst, dst)
	}
	err = ioutil.WriteFile(dst, []byte(staticContent), 0666)
	if err != nil {
		log.Fatal(err)
	}
}

func generatePkgMap(pkgs map[string]string, reservedIdents []string) string {
	var keys []string
	for k := range pkgs {
		keys = append(keys, k)
	}
	sort.Strings(keys)
	var buf bytes.Buffer
	buf.WriteString(`
	// pkgIdentMap maps from package name to one of the package's identifiers.
	// These identifiers are used by the Gorums protoc plugin to generate
	// appropriate import statements.
	`)
	buf.WriteString("var pkgIdentMap = map[string]string{\n")
	for _, imp := range keys {
		buf.WriteString("\t" + `"`)
		buf.WriteString(imp)
		buf.WriteString(`": "`)
		buf.WriteString(pkgs[imp])
		buf.WriteString(`",` + "\n")
	}
	buf.WriteString("}\n\n")
	buf.WriteString(`
	// reservedIdents holds the set of Gorums reserved identifiers.
	// These identifiers cannot be used to define message types in a proto file.
	`)
	buf.WriteString("var reservedIdents = []string{\n")
	for _, ident := range reservedIdents {
		buf.WriteString("\t" + `"`)
		buf.WriteString(ident)
		buf.WriteString(`",` + "\n")
	}
	buf.WriteString("}\n\n")
	return buf.String()
}

// findIdentifiers examines the given package to find all imported packages,
// and one used identifier in that imported package. These identifiers are
// used by the Gorums protoc plugin to generate appropriate import statements.
func findIdentifiers(fset *token.FileSet, info *loader.PackageInfo) (map[string]string, []string) {
	packagePath := info.Pkg.Path()
	pkgIdents := make(map[string][]string)
	for id, obj := range info.Uses {
		pos := fset.Position(id.Pos())
		if strings.Contains(pos.Filename, "zorums") {
			// ignore identifiers in zorums generated files
			continue
		}
		r := []rune(obj.Name())
		if unicode.IsLower(r[0]) {
			// ignore unexported identifiers
			continue
		}
		if pkg := obj.Pkg(); pkg != nil {
			switch obj := obj.(type) {
			case *types.Func:
				if typ := obj.Type(); typ != nil {
					if recv := typ.(*types.Signature).Recv(); recv != nil {
						// ignore functions on non-package types
						continue
					}
				}
				addUniqueIdentifier(pkgIdents, pkg.Path(), obj.Name())

			case *types.Const, *types.TypeName:
				addUniqueIdentifier(pkgIdents, pkg.Path(), obj.Name())

			case *types.Var:
				// no need to import obj's pkg if obj is a field of a struct.
				if obj.IsField() {
					continue
				}
				addUniqueIdentifier(pkgIdents, pkg.Path(), obj.Name())
			}
		}
	}
	// Only need to store one identifier for each imported package.
	// However, to ensure stable output, we sort the identifiers
	// and use the first element.
	pkgIdent := make(map[string]string, len(pkgIdents))
	var reservedIdents []string
	for path, idents := range pkgIdents {
		sort.Strings(idents)
		if path == packagePath {
			reservedIdents = idents
			continue
		}
		pkgIdent[path] = idents[0]
	}
	return pkgIdent, reservedIdents
}

func addUniqueIdentifier(pkgIdents map[string][]string, path, name string) {
	currentNames := pkgIdents[path]
	for _, known := range currentNames {
		if name == known {
			return
		}
	}
	pkgIdents[path] = append(pkgIdents[path], name)
}

// bundle returns a slice with the code for the given src package without imports.
// The returned map contains packages to be imported along with one identifier
// using the relevant import path.
// Loosly based on x/tools/cmd/bundle
func bundle(src string) (map[string]string, []string, []byte) {
	conf := loader.Config{ParserMode: parser.ParseComments, Build: &build.Default}
	conf.Import(src)
	lprog, err := conf.Load()
	if err != nil {
		log.Fatalf("failed to load Go package: %v", err)
	}
	// Because there was a single Import call and Load succeeded,
	// InitialPackages is guaranteed to hold the sole requested package.
	info := lprog.InitialPackages()[0]

	var out bytes.Buffer
	printFiles(&out, lprog.Fset, info)
	pkgIdentMap, reservedIdents := findIdentifiers(lprog.Fset, info)
	return pkgIdentMap, reservedIdents, out.Bytes()
}

func printFiles(out *bytes.Buffer, fset *token.FileSet, info *loader.PackageInfo) {
	for _, f := range info.Files {
		// filter files in dev package that shouldn't be bundled in template_static.go
		fileName := fset.File(f.Pos()).Name()
		if ignore(fileName) {
			continue
		}

		last := f.Package
		if len(f.Imports) > 0 {
			imp := f.Imports[len(f.Imports)-1]
			last = imp.End()
			if imp.Comment != nil {
				if e := imp.Comment.End(); e > last {
					last = e
				}
			}
		}

		// Pretty-print package-level declarations.
		// but no package or import declarations.
		var buf bytes.Buffer
		for _, decl := range f.Decls {
			if decl, ok := decl.(*ast.GenDecl); ok && decl.Tok == token.IMPORT {
				continue
			}
			beg, end := sourceRange(decl)
			printComments(out, f.Comments, last, beg)

			buf.Reset()
			format.Node(&buf, fset, &printer.CommentedNode{Node: decl, Comments: f.Comments})
			out.Write(buf.Bytes())
			last = printSameLineComment(out, f.Comments, fset, end)
			out.WriteString("\n\n")
		}
		printLastComments(out, f.Comments, last)
	}
}

// ignore files in dev folder with suffixes that shouldn't be bundled.
func ignore(file string) bool {
	for _, suffix := range []string{".proto", "_test.go"} {
		if strings.HasSuffix(file, suffix) {
			return true
		}
	}
	base := path.Base(file)
	// ignore zorums* files
	return strings.HasPrefix(base, "zorums")
}

// sourceRange returns the [beg, end) interval of source code
// belonging to decl (incl. associated comments).
func sourceRange(decl ast.Decl) (beg, end token.Pos) {
	beg = decl.Pos()
	end = decl.End()

	var doc, com *ast.CommentGroup

	switch d := decl.(type) {
	case *ast.GenDecl:
		doc = d.Doc
		if len(d.Specs) > 0 {
			switch spec := d.Specs[len(d.Specs)-1].(type) {
			case *ast.ValueSpec:
				com = spec.Comment
			case *ast.TypeSpec:
				com = spec.Comment
			}
		}
	case *ast.FuncDecl:
		doc = d.Doc
	}

	if doc != nil {
		beg = doc.Pos()
	}
	if com != nil && com.End() > end {
		end = com.End()
	}

	return beg, end
}

func printPackageComments(out *bytes.Buffer, files []*ast.File) {
	// Concatenate package comments from all files...
	for _, f := range files {
		if doc := f.Doc.Text(); strings.TrimSpace(doc) != "" {
			for _, line := range strings.Split(doc, "\n") {
				fmt.Fprintf(out, "// %s\n", line)
			}
		}
	}
	// ...but don't let them become the actual package comment.
	fmt.Fprintln(out)
}

func printComments(out *bytes.Buffer, comments []*ast.CommentGroup, pos, end token.Pos) {
	for _, cg := range comments {
		if pos <= cg.Pos() && cg.Pos() < end {
			for _, c := range cg.List {
				fmt.Fprintln(out, c.Text)
			}
			fmt.Fprintln(out)
		}
	}
}

const infinity = 1 << 30

func printLastComments(out *bytes.Buffer, comments []*ast.CommentGroup, pos token.Pos) {
	printComments(out, comments, pos, infinity)
}

func printSameLineComment(out *bytes.Buffer, comments []*ast.CommentGroup, fset *token.FileSet, pos token.Pos) token.Pos {
	tf := fset.File(pos)
	for _, cg := range comments {
		if pos <= cg.Pos() && tf.Line(cg.Pos()) == tf.Line(pos) {
			for _, c := range cg.List {
				fmt.Fprintln(out, c.Text)
			}
			return cg.End()
		}
	}
	return pos
}
<|endoftext|>


<commit_before>package schema2

import (
	"encoding/json"
	"errors"
	"fmt"

	"github.com/docker/distribution"
	"github.com/docker/distribution/digest"
	"github.com/docker/distribution/manifest"
)

const (
	// MediaTypeManifest specifies the mediaType for the current version.
	MediaTypeManifest = "application/vnd.docker.distribution.manifest.v2+json"

	// MediaTypeConfig specifies the mediaType for the image configuration.
	MediaTypeConfig = "application/vnd.docker.container.image.v1+json"

	// MediaTypePluginConfig specifies the mediaType for plugin configuration.
	MediaTypePluginConfig = "application/vnd.docker.plugin.v0+json"

	// MediaTypeLayer is the mediaType used for layers referenced by the
	// manifest.
	MediaTypeLayer = "application/vnd.docker.image.rootfs.diff.tar.gzip"

	// MediaTypeForeignLayer is the mediaType used for layers that must be
	// downloaded from foreign URLs.
	MediaTypeForeignLayer = "application/vnd.docker.image.rootfs.foreign.diff.tar.gzip"
)

var (
	// SchemaVersion provides a pre-initialized version structure for this
	// packages version of the manifest.
	SchemaVersion = manifest.Versioned{
		SchemaVersion: 2,
		MediaType:     MediaTypeManifest,
	}
)

func init() {
	schema2Func := func(b []byte) (distribution.Manifest, distribution.Descriptor, error) {
		m := new(DeserializedManifest)
		err := m.UnmarshalJSON(b)
		if err != nil {
			return nil, distribution.Descriptor{}, err
		}

		dgst := digest.FromBytes(b)
		return m, distribution.Descriptor{Digest: dgst, Size: int64(len(b)), MediaType: MediaTypeManifest}, err
	}
	err := distribution.RegisterManifestSchema(MediaTypeManifest, schema2Func)
	if err != nil {
		panic(fmt.Sprintf("Unable to register manifest: %s", err))
	}
}

// Manifest defines a schema2 manifest.
type Manifest struct {
	manifest.Versioned

	// Config references the image configuration as a blob.
	Config distribution.Descriptor `json:"config"`

	// Layers lists descriptors for the layers referenced by the
	// configuration.
	Layers []distribution.Descriptor `json:"layers"`
}

// References returnes the descriptors of this manifests references.
func (m Manifest) References() []distribution.Descriptor {
	return m.Layers
}

// Target returns the target of this signed manifest.
func (m Manifest) Target() distribution.Descriptor {
	return m.Config
}

// DeserializedManifest wraps Manifest with a copy of the original JSON.
// It satisfies the distribution.Manifest interface.
type DeserializedManifest struct {
	Manifest

	// canonical is the canonical byte representation of the Manifest.
	canonical []byte
}

// FromStruct takes a Manifest structure, marshals it to JSON, and returns a
// DeserializedManifest which contains the manifest and its JSON representation.
func FromStruct(m Manifest) (*DeserializedManifest, error) {
	var deserialized DeserializedManifest
	deserialized.Manifest = m

	var err error
	deserialized.canonical, err = json.MarshalIndent(&m, "", "   ")
	return &deserialized, err
}

// UnmarshalJSON populates a new Manifest struct from JSON data.
func (m *DeserializedManifest) UnmarshalJSON(b []byte) error {
	m.canonical = make([]byte, len(b), len(b))
	// store manifest in canonical
	copy(m.canonical, b)

	// Unmarshal canonical JSON into Manifest object
	var manifest Manifest
	if err := json.Unmarshal(m.canonical, &manifest); err != nil {
		return err
	}

	m.Manifest = manifest

	return nil
}

// MarshalJSON returns the contents of canonical. If canonical is empty,
// marshals the inner contents.
func (m *DeserializedManifest) MarshalJSON() ([]byte, error) {
	if len(m.canonical) > 0 {
		return m.canonical, nil
	}

	return nil, errors.New("JSON representation not initialized in DeserializedManifest")
}

// Payload returns the raw content of the manifest. The contents can be used to
// calculate the content identifier.
func (m DeserializedManifest) Payload() (string, []byte, error) {
	return m.MediaType, m.canonical, nil
}
<commit_msg>upate plugin MediaType<commit_after>package schema2

import (
	"encoding/json"
	"errors"
	"fmt"

	"github.com/docker/distribution"
	"github.com/docker/distribution/digest"
	"github.com/docker/distribution/manifest"
)

const (
	// MediaTypeManifest specifies the mediaType for the current version.
	MediaTypeManifest = "application/vnd.docker.distribution.manifest.v2+json"

	// MediaTypeConfig specifies the mediaType for the image configuration.
	MediaTypeConfig = "application/vnd.docker.container.image.v1+json"

	// MediaTypePluginConfig specifies the mediaType for plugin configuration.
	MediaTypePluginConfig = "application/vnd.docker.plugin.image.v0+json"

	// MediaTypeLayer is the mediaType used for layers referenced by the
	// manifest.
	MediaTypeLayer = "application/vnd.docker.image.rootfs.diff.tar.gzip"

	// MediaTypeForeignLayer is the mediaType used for layers that must be
	// downloaded from foreign URLs.
	MediaTypeForeignLayer = "application/vnd.docker.image.rootfs.foreign.diff.tar.gzip"
)

var (
	// SchemaVersion provides a pre-initialized version structure for this
	// packages version of the manifest.
	SchemaVersion = manifest.Versioned{
		SchemaVersion: 2,
		MediaType:     MediaTypeManifest,
	}
)

func init() {
	schema2Func := func(b []byte) (distribution.Manifest, distribution.Descriptor, error) {
		m := new(DeserializedManifest)
		err := m.UnmarshalJSON(b)
		if err != nil {
			return nil, distribution.Descriptor{}, err
		}

		dgst := digest.FromBytes(b)
		return m, distribution.Descriptor{Digest: dgst, Size: int64(len(b)), MediaType: MediaTypeManifest}, err
	}
	err := distribution.RegisterManifestSchema(MediaTypeManifest, schema2Func)
	if err != nil {
		panic(fmt.Sprintf("Unable to register manifest: %s", err))
	}
}

// Manifest defines a schema2 manifest.
type Manifest struct {
	manifest.Versioned

	// Config references the image configuration as a blob.
	Config distribution.Descriptor `json:"config"`

	// Layers lists descriptors for the layers referenced by the
	// configuration.
	Layers []distribution.Descriptor `json:"layers"`
}

// References returnes the descriptors of this manifests references.
func (m Manifest) References() []distribution.Descriptor {
	return m.Layers
}

// Target returns the target of this signed manifest.
func (m Manifest) Target() distribution.Descriptor {
	return m.Config
}

// DeserializedManifest wraps Manifest with a copy of the original JSON.
// It satisfies the distribution.Manifest interface.
type DeserializedManifest struct {
	Manifest

	// canonical is the canonical byte representation of the Manifest.
	canonical []byte
}

// FromStruct takes a Manifest structure, marshals it to JSON, and returns a
// DeserializedManifest which contains the manifest and its JSON representation.
func FromStruct(m Manifest) (*DeserializedManifest, error) {
	var deserialized DeserializedManifest
	deserialized.Manifest = m

	var err error
	deserialized.canonical, err = json.MarshalIndent(&m, "", "   ")
	return &deserialized, err
}

// UnmarshalJSON populates a new Manifest struct from JSON data.
func (m *DeserializedManifest) UnmarshalJSON(b []byte) error {
	m.canonical = make([]byte, len(b), len(b))
	// store manifest in canonical
	copy(m.canonical, b)

	// Unmarshal canonical JSON into Manifest object
	var manifest Manifest
	if err := json.Unmarshal(m.canonical, &manifest); err != nil {
		return err
	}

	m.Manifest = manifest

	return nil
}

// MarshalJSON returns the contents of canonical. If canonical is empty,
// marshals the inner contents.
func (m *DeserializedManifest) MarshalJSON() ([]byte, error) {
	if len(m.canonical) > 0 {
		return m.canonical, nil
	}

	return nil, errors.New("JSON representation not initialized in DeserializedManifest")
}

// Payload returns the raw content of the manifest. The contents can be used to
// calculate the content identifier.
func (m DeserializedManifest) Payload() (string, []byte, error) {
	return m.MediaType, m.canonical, nil
}
<|endoftext|>


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

package docker

import (
	"github.com/globocom/tsuru/cmd"
	"github.com/globocom/tsuru/exec/testing"
	"launchpad.net/gocheck"
	"net/http"
	"net/http/httptest"
	"strings"
)

var _ = gocheck.Suite(SSHSuite{})

type SSHSuite struct{}

func (SSHSuite) TestExecuteCommandHandler(c *gocheck.C) {
	output := []byte(". ..")
	out := map[string][][]byte{"*": {output}}
	fexec := &testing.FakeExecutor{Output: out}
	setExecut(fexec)
	defer setExecut(nil)
	data := `{"cmd":"ls","args":["-l", "-a"]}`
	body := strings.NewReader(data)
	request, _ := http.NewRequest("POST", "/container/10.10.10.10/cmd?:ip=10.10.10.10", body)
	request.Header.Set("Content-Type", "application/json")
	recorder := httptest.NewRecorder()
	sshHandler(recorder, request)
	c.Assert(recorder.Code, gocheck.Equals, http.StatusOK)
	c.Assert(recorder.Body.String(), gocheck.Equals, string(output))
	c.Assert(recorder.Header().Get("Content-Type"), gocheck.Equals, "text")
	args := []string{
		"10.10.10.10", "-l", "ubuntu",
		"-o", "StrictHostKeyChecking no",
		"--", "ls", "-l", "-a",
	}
	c.Assert(fexec.ExecutedCmd("ssh", args), gocheck.Equals, true)
}

func (SSHSuite) TestExecuteCommandHandlerInvalidJSON(c *gocheck.C) {
	data := `}}}}---;"`
	body := strings.NewReader(data)
	request, _ := http.NewRequest("POST", "/container/10.10.10.10/cmd?:ip=10.10.10.10", body)
	request.Header.Set("Content-Type", "application/json")
	recorder := httptest.NewRecorder()
	sshHandler(recorder, request)
	c.Assert(recorder.Code, gocheck.Equals, http.StatusBadRequest)
	c.Assert(recorder.Body.String(), gocheck.Equals, "Invalid JSON\n")
}

func (SSHSuite) TestSSHAgentCmdInfo(c *gocheck.C) {
	desc := `Start HTTP agent for running commands on Docker via SSH.

By default, the agent will listen on 0.0.0.0:4545. Use --listen or -l to
specify the address to listen on.
`
	expected := &cmd.Info{
		Name:  "docker-ssh-agent",
		Usage: "docker-ssh-agent",
		Desc:  desc,
	}
	c.Assert((&sshAgentCmd{}).Info(), gocheck.DeepEquals, expected)
}

func (SSHSuite) TestSSHAgentCmdFlags(c *gocheck.C) {
	cmd := sshAgentCmd{}
	flags := cmd.Flags()
	flags.Parse(true, []string{})
	flag := flags.Lookup("l")
	c.Check(flag.Name, gocheck.Equals, "l")
	c.Check(flag.DefValue, gocheck.Equals, "0.0.0.0:4545")
	c.Check(flag.Usage, gocheck.Equals, "Address to listen on")
	flag = flags.Lookup("listen")
	c.Check(flag.Name, gocheck.Equals, "listen")
	c.Check(flag.DefValue, gocheck.Equals, "0.0.0.0:4545")
	c.Check(flag.Usage, gocheck.Equals, "Address to listen on")
	c.Check(cmd.listen, gocheck.Equals, "0.0.0.0:4545")
}
<commit_msg>provision/docker: add FakeSSHServer<commit_after>// Copyright 2013 tsuru authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package docker

import (
	"encoding/json"
	"github.com/globocom/tsuru/cmd"
	"github.com/globocom/tsuru/exec/testing"
	"launchpad.net/gocheck"
	"net/http"
	"net/http/httptest"
	"strings"
)

var _ = gocheck.Suite(SSHSuite{})

type SSHSuite struct{}

func (SSHSuite) TestExecuteCommandHandler(c *gocheck.C) {
	output := []byte(". ..")
	out := map[string][][]byte{"*": {output}}
	fexec := &testing.FakeExecutor{Output: out}
	setExecut(fexec)
	defer setExecut(nil)
	data := `{"cmd":"ls","args":["-l", "-a"]}`
	body := strings.NewReader(data)
	request, _ := http.NewRequest("POST", "/container/10.10.10.10/cmd?:ip=10.10.10.10", body)
	request.Header.Set("Content-Type", "application/json")
	recorder := httptest.NewRecorder()
	sshHandler(recorder, request)
	c.Assert(recorder.Code, gocheck.Equals, http.StatusOK)
	c.Assert(recorder.Body.String(), gocheck.Equals, string(output))
	c.Assert(recorder.Header().Get("Content-Type"), gocheck.Equals, "text")
	args := []string{
		"10.10.10.10", "-l", "ubuntu",
		"-o", "StrictHostKeyChecking no",
		"--", "ls", "-l", "-a",
	}
	c.Assert(fexec.ExecutedCmd("ssh", args), gocheck.Equals, true)
}

func (SSHSuite) TestExecuteCommandHandlerInvalidJSON(c *gocheck.C) {
	data := `}}}}---;"`
	body := strings.NewReader(data)
	request, _ := http.NewRequest("POST", "/container/10.10.10.10/cmd?:ip=10.10.10.10", body)
	request.Header.Set("Content-Type", "application/json")
	recorder := httptest.NewRecorder()
	sshHandler(recorder, request)
	c.Assert(recorder.Code, gocheck.Equals, http.StatusBadRequest)
	c.Assert(recorder.Body.String(), gocheck.Equals, "Invalid JSON\n")
}

func (SSHSuite) TestSSHAgentCmdInfo(c *gocheck.C) {
	desc := `Start HTTP agent for running commands on Docker via SSH.

By default, the agent will listen on 0.0.0.0:4545. Use --listen or -l to
specify the address to listen on.
`
	expected := &cmd.Info{
		Name:  "docker-ssh-agent",
		Usage: "docker-ssh-agent",
		Desc:  desc,
	}
	c.Assert((&sshAgentCmd{}).Info(), gocheck.DeepEquals, expected)
}

func (SSHSuite) TestSSHAgentCmdFlags(c *gocheck.C) {
	cmd := sshAgentCmd{}
	flags := cmd.Flags()
	flags.Parse(true, []string{})
	flag := flags.Lookup("l")
	c.Check(flag.Name, gocheck.Equals, "l")
	c.Check(flag.DefValue, gocheck.Equals, "0.0.0.0:4545")
	c.Check(flag.Usage, gocheck.Equals, "Address to listen on")
	flag = flags.Lookup("listen")
	c.Check(flag.Name, gocheck.Equals, "listen")
	c.Check(flag.DefValue, gocheck.Equals, "0.0.0.0:4545")
	c.Check(flag.Usage, gocheck.Equals, "Address to listen on")
	c.Check(cmd.listen, gocheck.Equals, "0.0.0.0:4545")
}

type FakeSSHServer struct {
	requests []*http.Request
	bodies   []cmdInput
	output   string
}

func (h *FakeSSHServer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	var input cmdInput
	defer r.Body.Close()
	err := json.NewDecoder(r.Body).Decode(&input)
	if err != nil {
		panic(err)
	}
	h.requests = append(h.requests, r)
	h.bodies = append(h.bodies, input)
	w.Write([]byte(h.output))
}
<|endoftext|>


<commit_before>package proxy

import (
	"net"
	"net/http"

	"github.com/cloudfoundry/gorouter/route"
)

type ProxyRoundTripper struct {
	Transport      http.RoundTripper
	After          AfterRoundTrip
	Iter           route.EndpointIterator
	Handler        *RequestHandler
	ServingBackend bool
}

type RoundTripEventHandler interface {
	BeforeRoundTrip(request *http.Request) (*route.Endpoint, error)
	HandleError(err error)
}

type BackendRoundTripper struct {
	Iter    route.EndpointIterator
	Handler *RequestHandler
}

type RouteServiceRoundTripper struct {
	Handler *RequestHandler
}

func (be *BackendRoundTripper) BeforeRoundTrip(request *http.Request) (*route.Endpoint, error) {
	endpoint := be.Iter.Next()

	if endpoint == nil {
		be.Handler.reporter.CaptureBadGateway(request)
		err := noEndpointsAvailable
		be.Handler.HandleBadGateway(err)
		return nil, err
	}
	be.setupBackendRequest(request, endpoint)

	return endpoint, nil
}

func (be *BackendRoundTripper) HandleError(err error) {
	be.Iter.EndpointFailed()
	be.Handler.Logger().Set("Error", err.Error())
	be.Handler.Logger().Warnf("proxy.endpoint.failed")
}

func (rs *RouteServiceRoundTripper) BeforeRoundTrip(request *http.Request) (*route.Endpoint, error) {
	return newRouteServiceEndpoint(), nil
}

func (rs *RouteServiceRoundTripper) HandleError(err error) {
	rs.Handler.Logger().Set("Error", err.Error())
	rs.Handler.Logger().Warnf("proxy.route-service.failed")
}

func (be *BackendRoundTripper) setupBackendRequest(request *http.Request, endpoint *route.Endpoint) {
	be.Handler.Logger().Debug("proxy.backend")

	request.URL.Host = endpoint.CanonicalAddr()
	request.Header.Set("X-CF-ApplicationID", endpoint.ApplicationId)
	setRequestXCfInstanceId(request, endpoint)
}

func (p *ProxyRoundTripper) RoundTrip(request *http.Request) (*http.Response, error) {
	var err error
	var res *http.Response
	var endpoint *route.Endpoint

	handler := p.eventHandler()

	retry := 0
	for {
		endpoint, err = handler.BeforeRoundTrip(request)
		if err != nil {
			return nil, err
		}

		res, err = p.Transport.RoundTrip(request)

		if err == nil {
			break
		}
		if ne, netErr := err.(*net.OpError); !netErr || ne.Op != "dial" {
			break
		}

		handler.HandleError(err)

		retry++
		if retry == retries {
			break
		}
	}

	if p.After != nil {
		p.After(res, endpoint, err)
	}

	return res, err
}

func (p *ProxyRoundTripper) eventHandler() RoundTripEventHandler {
	var r RoundTripEventHandler
	if p.ServingBackend {
		r = &BackendRoundTripper{
			Iter:    p.Iter,
			Handler: p.Handler,
		}
	} else {
		r = &RouteServiceRoundTripper{
			Handler: p.Handler,
		}
	}
	return r
}
<commit_msg>Refactor RoundTripEventHandler<commit_after>package proxy

import (
	"net"
	"net/http"

	"github.com/cloudfoundry/gorouter/route"
)

type ProxyRoundTripper struct {
	Transport      http.RoundTripper
	After          AfterRoundTrip
	Iter           route.EndpointIterator
	Handler        *RequestHandler
	ServingBackend bool
}

type RoundTripEventHandler interface {
	SelectEndpoint(request *http.Request) (*route.Endpoint, error)
	PreprocessRequest(request *http.Request, endpoint *route.Endpoint)
	HandleError(err error)
}

type BackendRoundTripper struct {
	Iter    route.EndpointIterator
	Handler *RequestHandler
}

type RouteServiceRoundTripper struct {
	Handler *RequestHandler
}

func (be *BackendRoundTripper) SelectEndpoint(request *http.Request) (*route.Endpoint, error) {
	endpoint := be.Iter.Next()

	if endpoint == nil {
		be.Handler.reporter.CaptureBadGateway(request)
		err := noEndpointsAvailable
		be.Handler.HandleBadGateway(err)
		return nil, err
	}
	return endpoint, nil
}

func (be *BackendRoundTripper) HandleError(err error) {
	be.Iter.EndpointFailed()
	be.Handler.Logger().Set("Error", err.Error())
	be.Handler.Logger().Warnf("proxy.endpoint.failed")
}

func (rs *RouteServiceRoundTripper) SelectEndpoint(request *http.Request) (*route.Endpoint, error) {
	return newRouteServiceEndpoint(), nil
}
func (be *RouteServiceRoundTripper) PreprocessRequest(request *http.Request, endpoint *route.Endpoint) {
}

func (rs *RouteServiceRoundTripper) HandleError(err error) {
	rs.Handler.Logger().Set("Error", err.Error())
	rs.Handler.Logger().Warnf("proxy.route-service.failed")
}

func (be *BackendRoundTripper) PreprocessRequest(request *http.Request, endpoint *route.Endpoint) {
	be.Handler.Logger().Debug("proxy.backend")

	request.URL.Host = endpoint.CanonicalAddr()
	request.Header.Set("X-CF-ApplicationID", endpoint.ApplicationId)
	setRequestXCfInstanceId(request, endpoint)
}

func (p *ProxyRoundTripper) RoundTrip(request *http.Request) (*http.Response, error) {
	var err error
	var res *http.Response
	var endpoint *route.Endpoint

	handler := p.eventHandler()

	retry := 0
	for {
		endpoint, err = handler.SelectEndpoint(request)
		if err != nil {
			return nil, err
		}

		handler.PreprocessRequest(request, endpoint)

		res, err = p.Transport.RoundTrip(request)
		if err == nil {
			break
		}
		if ne, netErr := err.(*net.OpError); !netErr || ne.Op != "dial" {
			break
		}

		handler.HandleError(err)

		retry++
		if retry == retries {
			break
		}
	}

	if p.After != nil {
		p.After(res, endpoint, err)
	}

	return res, err
}

func (p *ProxyRoundTripper) eventHandler() RoundTripEventHandler {
	var r RoundTripEventHandler
	if p.ServingBackend {
		r = &BackendRoundTripper{
			Iter:    p.Iter,
			Handler: p.Handler,
		}
	} else {
		r = &RouteServiceRoundTripper{
			Handler: p.Handler,
		}
	}
	return r
}
<|endoftext|>


<commit_before>package statemachine_test

import (
	"testing"
	"time"

	"github.com/lytics/metafora"
	"github.com/lytics/metafora/embedded"
	"github.com/lytics/metafora/statemachine"
)

func testhandler(tid string, cmds <-chan statemachine.Message) statemachine.Message {
	metafora.Debugf("Starting %s", tid)
	m := <-cmds
	metafora.Debugf("%s recvd %s", tid, m.Code)
	return m
}

type testStore struct {
	initial *statemachine.State
	out     chan<- *statemachine.State
}

func (s testStore) Load(taskID string) (*statemachine.State, error) {
	s.out <- s.initial
	return s.initial, nil
}
func (s testStore) Store(taskID string, newstate *statemachine.State) error {
	metafora.Debugf("%s storing %s", taskID, newstate.Code)
	s.out <- newstate
	return nil
}

// setup a task with the specified task ID in a stateful handler and run it.
func setup(t *testing.T, tid string) (*embedded.StateStore, statemachine.Commander, metafora.Handler, chan bool) {
	t.Parallel()
	ss := embedded.NewStateStore().(*embedded.StateStore)
	ss.Store(tid, &statemachine.State{Code: statemachine.Runnable})
	<-ss.Stored // pop initial state out
	cmdr := embedded.NewCommander()
	cmdlistener := cmdr.NewListener(tid)
	sm := statemachine.New(tid, testhandler, ss, cmdlistener, nil)
	done := make(chan bool)
	go func() { done <- sm.Run() }()
	return ss, cmdr, sm, done
}

//FIXME leaks goroutines
func TestRules(t *testing.T) {
	t.Parallel()
	for i, trans := range statemachine.Rules {
		metafora.Debugf("Trying %s", trans)
		cmdr := embedded.NewCommander()
		cmdlistener := cmdr.NewListener("test")
		store := make(chan *statemachine.State)

		state := &statemachine.State{Code: trans.From}

		// Sleeping state needs extra Until state
		if trans.From == statemachine.Sleeping {
			until := time.Now().Add(100 * time.Millisecond)
			state.Until = &until
		}

		ts := testStore{initial: state, out: store}

		// Create a new statemachine that starts from the From state
		sm := statemachine.New("test", testhandler, ts, cmdlistener, nil)
		go sm.Run()
		initial := <-store
		if initial.Code != trans.From {
			t.Fatalf("%d Initial state %q not set. Found: %q", i, trans.From, initial.Code)
		}

		// The Fault state transitions itself to either sleeping or failed
		if trans.From != statemachine.Fault {
			// Apply the Event to transition to the To state
			msg := statemachine.Message{Code: trans.Event}

			// Sleep messages need extra state
			if trans.Event == statemachine.Sleep {
				until := time.Now().Add(10 * time.Millisecond)
				msg.Until = &until
			}
			if err := cmdr.Send("test", statemachine.Message{Code: trans.Event}); err != nil {
				t.Fatalf("Error sending message %s: %v", trans.Event, err)
			}
		}
		newstate := <-store
		if trans.From == statemachine.Fault && trans.To == statemachine.Failed {
			// continue on as this transition relies on state this test doesn't exercise
			continue
		}
		if newstate.Code != trans.To {
			t.Fatalf("%d Expected %q but found %q for transition %s", i, trans.To, newstate.Code, trans)
		}
	}
}

func TestCheckpointRelease(t *testing.T) {
	ss, cmdr, _, done := setup(t, "test1")

	// Should just cause statemachine to loop
	if err := cmdr.Send("test1", statemachine.Message{Code: statemachine.Checkpoint}); err != nil {
		t.Fatalf("Error sending checkpoint: %v", err)
	}
	select {
	case <-done:
		t.Fatalf("Checkpoint command should not have caused statemachine to exit.")
	case <-time.After(100 * time.Millisecond):
	}

	// Should cause the statemachine to exit
	if err := cmdr.Send("test1", statemachine.Message{Code: statemachine.Release}); err != nil {
		t.Fatalf("Error sending release: %v", err)
	}
	select {
	case d := <-done:
		if d {
			t.Fatalf("Release command should not have caused the task to be marked as done.")
		}
	case <-time.After(100 * time.Millisecond):
		t.Fatalf("Expected statemachine to exit but it did not.")
	}
	state, err := ss.Load("test1")
	if err != nil {
		t.Fatal(err)
	}
	if state.Code != statemachine.Runnable {
		t.Fatalf("Expected released task to be runnable but found state %q", state.Code)
	}
}

func TestSleep(t *testing.T) {
	ss, cmdr, _, _ := setup(t, "sleep-test")

	{
		// Put to sleep forever
		until := time.Now().Add(9001 * time.Hour)
		if err := cmdr.Send("sleep-test", statemachine.Message{Code: statemachine.Sleep, Until: &until}); err != nil {
			t.Fatalf("Error sending sleep: %v", err)
		}

		newstate := <-ss.Stored
		if newstate.State.Code != statemachine.Sleeping || !newstate.State.Until.Equal(until) {
			t.Fatalf("Expected task to store state Sleeping, but stored: %s", newstate)
		}
	}

	// Make sure it stays sleeping for at least a bit
	select {
	case newstate := <-ss.Stored:
		t.Fatalf("Expected task to stay asleep forever but transitioned to: %s", newstate)
	case <-time.After(100 * time.Millisecond):
	}

	// Override current sleep with a shorter one
	dur := 1 * time.Second
	start := time.Now()
	until := start.Add(dur)
	if err := cmdr.Send("sleep-test", statemachine.Message{Code: statemachine.Sleep, Until: &until}); err != nil {
		t.Fatalf("Error sending sleep: %v", err)
	}

	newstate := <-ss.Stored
	if newstate.State.Code != statemachine.Sleeping || !newstate.State.Until.Equal(until) {
		t.Fatalf("Expected task to store state Sleeping, but stored: %s", newstate)
	}

	// Make sure it transitions to Runnable after sleep has elapsed
	newstate = <-ss.Stored
	transitioned := time.Now()
	if newstate.State.Code != statemachine.Runnable || newstate.State.Until != nil {
		t.Fatalf("Expected task to be runnable without an Until time but found: %s", newstate.State)
	}
	elapsed := transitioned.Sub(start)
	if transitioned.Sub(start) < dur {
		t.Fatalf("Expected task to sleep for %s but slept for %s", dur, elapsed)
	}
	t.Logf("Statemachine latency: %s", elapsed-dur)
}

func TestTerminal(t *testing.T) {
	ss, cmdr, sm, done := setup(t, "terminal-test")

	// Kill the task
	if err := cmdr.Send("terminal-test", statemachine.Message{Code: statemachine.Kill}); err != nil {
		t.Fatalf("Error sending kill command: %v", err)
	}

	// Task should be killed and done (unscheduled)
	newstate := <-ss.Stored
	if newstate.State.Code != statemachine.Killed {
		t.Fatalf("Expected task to be killed but found: %s", newstate.State)
	}
	if !(<-done) {
		t.Fatal("Expected task to be done.")
	}
	if state, err := ss.Load("terminal-test"); err != nil || state.Code != statemachine.Killed {
		t.Fatalf("Failed to load expected killed state for task: state=%s err=%v", state, err)
	}

	// Task should just die again if we try to reschedule it
	go func() { done <- sm.Run() }()
	select {
	case newstate := <-ss.Stored:
		t.Fatalf("Re-running a terminated task should *not* store state, but it stored: %v", newstate.State)
	case <-time.After(100 * time.Millisecond):
		// State shouldn't even be stored since it's not being changed and terminal
		// states should be immutable
	}

	if !(<-done) {
		t.Fatal("Expected task to be done.")
	}
}

func TestPause(t *testing.T) {
	ss, cmdr, _, done := setup(t, "test-pause")

	pause := func() {
		if err := cmdr.Send("test-pause", statemachine.Message{Code: statemachine.Pause}); err != nil {
			t.Fatalf("Error sending pause command to test-pause: %v", err)
		}
		newstate := <-ss.Stored
		if newstate.State.Code != statemachine.Paused {
			t.Fatalf("Expected paused state but found: %s", newstate.State)
		}
		if state, err := ss.Load("test-pause"); err != nil || state.Code != statemachine.Paused {
			t.Fatalf("Failed to load expected pause state for task: state=%s err=%v", state, err)
		}

		// Task should not be Done; pausing doesn't exit the statemachine
		select {
		case <-done:
			t.Fatal("Task exited unexpectedly.")
		case <-time.After(100 * time.Millisecond):
		}
	}

	// Pause the work
	pause()

	// Should be able to resume paused work
	if err := cmdr.Send("test-pause", statemachine.Message{Code: statemachine.Run}); err != nil {
		t.Fatalf("Error sending run command to test-pause: %v", err)
	}
	newstate := <-ss.Stored
	if newstate.State.Code != statemachine.Runnable {
		t.Fatalf("Expected runnable state but found: %s", newstate.State)
	}
	if state, err := ss.Load("test-pause"); err != nil || state.Code != statemachine.Runnable {
		t.Fatalf("Failed to load expected runnable state for task: state=%s err=%v", state, err)
	}

	// Re-pause the work
	pause()

	// Pausing paused work is silly but fine
	pause()

	// Releasing paused work should make it exit but leave it in the paused state
	if err := cmdr.Send("test-pause", statemachine.Message{Code: statemachine.Release}); err != nil {
		t.Fatalf("Error sending release command to test-pause: %v", err)
	}
	newstate = <-ss.Stored
	if newstate.State.Code != statemachine.Paused {
		t.Fatalf("Releasing should not have changed paused state but stored: %s", newstate.State)
	}
	select {
	case d := <-done:
		if d {
			t.Fatal("Releasing task should not have marked it as done.")
		}
	case <-time.After(100 * time.Millisecond):
		t.Fatal("Releasing paused task should have exited the statemachine, but didn't.")
	}

	// Ensure task is stored with the paused state
	if state, err := ss.Load("test-pause"); err != nil || state.Code != statemachine.Paused {
		t.Fatalf("Failed to load expected paused state for task: state=%s err=%v", state, err)
	}
}
<commit_msg>Use sm.Stop in a test as opposed to commands<commit_after>package statemachine_test

import (
	"testing"
	"time"

	"github.com/lytics/metafora"
	"github.com/lytics/metafora/embedded"
	"github.com/lytics/metafora/statemachine"
)

func testhandler(tid string, cmds <-chan statemachine.Message) statemachine.Message {
	metafora.Debugf("Starting %s", tid)
	m := <-cmds
	metafora.Debugf("%s recvd %s", tid, m.Code)
	return m
}

type testStore struct {
	initial *statemachine.State
	out     chan<- *statemachine.State
}

func (s testStore) Load(taskID string) (*statemachine.State, error) {
	s.out <- s.initial
	return s.initial, nil
}
func (s testStore) Store(taskID string, newstate *statemachine.State) error {
	metafora.Debugf("%s storing %s", taskID, newstate.Code)
	s.out <- newstate
	return nil
}

// setup a task with the specified task ID in a stateful handler and run it.
func setup(t *testing.T, tid string) (*embedded.StateStore, statemachine.Commander, metafora.Handler, chan bool) {
	t.Parallel()
	ss := embedded.NewStateStore().(*embedded.StateStore)
	ss.Store(tid, &statemachine.State{Code: statemachine.Runnable})
	<-ss.Stored // pop initial state out
	cmdr := embedded.NewCommander()
	cmdlistener := cmdr.NewListener(tid)
	sm := statemachine.New(tid, testhandler, ss, cmdlistener, nil)
	done := make(chan bool)
	go func() { done <- sm.Run() }()
	return ss, cmdr, sm, done
}

//FIXME leaks goroutines
func TestRules(t *testing.T) {
	t.Parallel()
	for i, trans := range statemachine.Rules {
		metafora.Debugf("Trying %s", trans)
		cmdr := embedded.NewCommander()
		cmdlistener := cmdr.NewListener("test")
		store := make(chan *statemachine.State)

		state := &statemachine.State{Code: trans.From}

		// Sleeping state needs extra Until state
		if trans.From == statemachine.Sleeping {
			until := time.Now().Add(100 * time.Millisecond)
			state.Until = &until
		}

		ts := testStore{initial: state, out: store}

		// Create a new statemachine that starts from the From state
		sm := statemachine.New("test", testhandler, ts, cmdlistener, nil)
		go sm.Run()
		initial := <-store
		if initial.Code != trans.From {
			t.Fatalf("%d Initial state %q not set. Found: %q", i, trans.From, initial.Code)
		}

		// The Fault state transitions itself to either sleeping or failed
		if trans.From != statemachine.Fault {
			// Apply the Event to transition to the To state
			msg := statemachine.Message{Code: trans.Event}

			// Sleep messages need extra state
			if trans.Event == statemachine.Sleep {
				until := time.Now().Add(10 * time.Millisecond)
				msg.Until = &until
			}
			if err := cmdr.Send("test", statemachine.Message{Code: trans.Event}); err != nil {
				t.Fatalf("Error sending message %s: %v", trans.Event, err)
			}
		}
		newstate := <-store
		if trans.From == statemachine.Fault && trans.To == statemachine.Failed {
			// continue on as this transition relies on state this test doesn't exercise
			continue
		}
		if newstate.Code != trans.To {
			t.Fatalf("%d Expected %q but found %q for transition %s", i, trans.To, newstate.Code, trans)
		}
	}
}

func TestCheckpointRelease(t *testing.T) {
	ss, cmdr, _, done := setup(t, "test1")

	// Should just cause statemachine to loop
	if err := cmdr.Send("test1", statemachine.Message{Code: statemachine.Checkpoint}); err != nil {
		t.Fatalf("Error sending checkpoint: %v", err)
	}
	select {
	case <-done:
		t.Fatalf("Checkpoint command should not have caused statemachine to exit.")
	case <-time.After(100 * time.Millisecond):
	}

	// Should cause the statemachine to exit
	if err := cmdr.Send("test1", statemachine.Message{Code: statemachine.Release}); err != nil {
		t.Fatalf("Error sending release: %v", err)
	}
	select {
	case d := <-done:
		if d {
			t.Fatalf("Release command should not have caused the task to be marked as done.")
		}
	case <-time.After(100 * time.Millisecond):
		t.Fatalf("Expected statemachine to exit but it did not.")
	}
	state, err := ss.Load("test1")
	if err != nil {
		t.Fatal(err)
	}
	if state.Code != statemachine.Runnable {
		t.Fatalf("Expected released task to be runnable but found state %q", state.Code)
	}
}

func TestSleep(t *testing.T) {
	ss, cmdr, _, _ := setup(t, "sleep-test")

	{
		// Put to sleep forever
		until := time.Now().Add(9001 * time.Hour)
		if err := cmdr.Send("sleep-test", statemachine.Message{Code: statemachine.Sleep, Until: &until}); err != nil {
			t.Fatalf("Error sending sleep: %v", err)
		}

		newstate := <-ss.Stored
		if newstate.State.Code != statemachine.Sleeping || !newstate.State.Until.Equal(until) {
			t.Fatalf("Expected task to store state Sleeping, but stored: %s", newstate)
		}
	}

	// Make sure it stays sleeping for at least a bit
	select {
	case newstate := <-ss.Stored:
		t.Fatalf("Expected task to stay asleep forever but transitioned to: %s", newstate)
	case <-time.After(100 * time.Millisecond):
	}

	// Override current sleep with a shorter one
	dur := 1 * time.Second
	start := time.Now()
	until := start.Add(dur)
	if err := cmdr.Send("sleep-test", statemachine.Message{Code: statemachine.Sleep, Until: &until}); err != nil {
		t.Fatalf("Error sending sleep: %v", err)
	}

	newstate := <-ss.Stored
	if newstate.State.Code != statemachine.Sleeping || !newstate.State.Until.Equal(until) {
		t.Fatalf("Expected task to store state Sleeping, but stored: %s", newstate)
	}

	// Make sure it transitions to Runnable after sleep has elapsed
	newstate = <-ss.Stored
	transitioned := time.Now()
	if newstate.State.Code != statemachine.Runnable || newstate.State.Until != nil {
		t.Fatalf("Expected task to be runnable without an Until time but found: %s", newstate.State)
	}
	elapsed := transitioned.Sub(start)
	if transitioned.Sub(start) < dur {
		t.Fatalf("Expected task to sleep for %s but slept for %s", dur, elapsed)
	}
	t.Logf("Statemachine latency: %s", elapsed-dur)
}

func TestTerminal(t *testing.T) {
	ss, cmdr, sm, done := setup(t, "terminal-test")

	// Kill the task
	if err := cmdr.Send("terminal-test", statemachine.Message{Code: statemachine.Kill}); err != nil {
		t.Fatalf("Error sending kill command: %v", err)
	}

	// Task should be killed and done (unscheduled)
	newstate := <-ss.Stored
	if newstate.State.Code != statemachine.Killed {
		t.Fatalf("Expected task to be killed but found: %s", newstate.State)
	}
	if !(<-done) {
		t.Fatal("Expected task to be done.")
	}
	if state, err := ss.Load("terminal-test"); err != nil || state.Code != statemachine.Killed {
		t.Fatalf("Failed to load expected killed state for task: state=%s err=%v", state, err)
	}

	// Task should just die again if we try to reschedule it
	go func() { done <- sm.Run() }()
	select {
	case newstate := <-ss.Stored:
		t.Fatalf("Re-running a terminated task should *not* store state, but it stored: %v", newstate.State)
	case <-time.After(100 * time.Millisecond):
		// State shouldn't even be stored since it's not being changed and terminal
		// states should be immutable
	}

	if !(<-done) {
		t.Fatal("Expected task to be done.")
	}
}

func TestPause(t *testing.T) {
	ss, cmdr, sm, done := setup(t, "test-pause")

	pause := func() {
		if err := cmdr.Send("test-pause", statemachine.Message{Code: statemachine.Pause}); err != nil {
			t.Fatalf("Error sending pause command to test-pause: %v", err)
		}
		newstate := <-ss.Stored
		if newstate.State.Code != statemachine.Paused {
			t.Fatalf("Expected paused state but found: %s", newstate.State)
		}
		if state, err := ss.Load("test-pause"); err != nil || state.Code != statemachine.Paused {
			t.Fatalf("Failed to load expected pause state for task: state=%s err=%v", state, err)
		}

		// Task should not be Done; pausing doesn't exit the statemachine
		select {
		case <-done:
			t.Fatal("Task exited unexpectedly.")
		case <-time.After(100 * time.Millisecond):
		}
	}

	// Pause the work
	pause()

	// Should be able to resume paused work
	if err := cmdr.Send("test-pause", statemachine.Message{Code: statemachine.Run}); err != nil {
		t.Fatalf("Error sending run command to test-pause: %v", err)
	}
	newstate := <-ss.Stored
	if newstate.State.Code != statemachine.Runnable {
		t.Fatalf("Expected runnable state but found: %s", newstate.State)
	}
	if state, err := ss.Load("test-pause"); err != nil || state.Code != statemachine.Runnable {
		t.Fatalf("Failed to load expected runnable state for task: state=%s err=%v", state, err)
	}

	// Re-pause the work
	pause()

	// Pausing paused work is silly but fine
	pause()

	// Releasing paused work should make it exit but leave it in the paused state
	sm.Stop()
	newstate = <-ss.Stored
	if newstate.State.Code != statemachine.Paused {
		t.Fatalf("Releasing should not have changed paused state but stored: %s", newstate.State)
	}
	select {
	case d := <-done:
		if d {
			t.Fatal("Releasing task should not have marked it as done.")
		}
	case <-time.After(100 * time.Millisecond):
		t.Fatal("Releasing paused task should have exited the statemachine, but didn't.")
	}

	// Ensure task is stored with the paused state
	if state, err := ss.Load("test-pause"); err != nil || state.Code != statemachine.Paused {
		t.Fatalf("Failed to load expected paused state for task: state=%s err=%v", state, err)
	}
}
<|endoftext|>


<commit_before>package authsym

import (
	"bytes"
	"fmt"
	"testing"
)

const (
	testEnc = "/tmp/test.out"
	testOut = "/tmp/test.dat"
	testRef = "testdata/vector01.dat"
	rPadRef = "testdata/vector02.dat"
)

var (
	testKey []byte
)

// FailWithError is a utility for dumping errors and failing the test.
func FailWithError(t *testing.T, err error) {
	fmt.Println("failed")
	if err != nil {
		fmt.Println("[!] ", err.Error())
	}
	t.FailNow()
}

// Test session key generation.
func TestGenerateKey(t *testing.T) {
	key, err := GenerateKey()
	if err != nil || len(key) != KeySize {
		FailWithError(t, err)
	}
}

// Test long term key generation.
/*
func TestGenerateLTKey(t *testing.T) {
	key, err := GenerateLTKey()
	if err != nil || len(key) != KeySize {
		FailWithError(t, err)
	}
}
 */

// Test initialisation vector generation.
func TestGenerateIV(t *testing.T) {
	iv, err := GenerateIV()
	if err != nil || len(iv) != BlockSize {
		FailWithError(t, err)
	}
}

// Does D(E(k,m)) == m?
func TestEncryptDecryptBlock(t *testing.T) {
	fmt.Println("encrypt block")
	m := []byte("Hello, world.")
	key, err := GenerateKey()
	if err != nil {
		FailWithError(t, err)
	}

	fmt.Println("\tencrypt")
	e, err := Encrypt(key, m)
	if err != nil {
		FailWithError(t, err)
	}

	fmt.Println("\tdecrypt")
	decrypted, err := Decrypt(key, e)
	if err != nil {
		FailWithError(t, err)
	}

	if !bytes.Equal(decrypted, m) {
		err = fmt.Errorf("plaintext doesn't match original message")
		FailWithError(t, err)
	}
	fmt.Println("finish encrypt block")
}

func TestEncryptDecryptBlockFails(t *testing.T) {
	m := []byte("Hello, world.")

	key, err := GenerateKey()
	if err != nil {
		FailWithError(t, err)
	}

	e, err := Encrypt(key, m)
	if err != nil {
		FailWithError(t, err)
	}

	n := len(e) - 2
	orig := e[n]
	if e[n] == 255 {
		e[n] = 0
	} else {
		e[n]++
	}
	if e[n] == orig {
		err = fmt.Errorf("byte not modified")
		FailWithError(t, err)
	}
	decrypted, err := Decrypt(key, e)
	if err == nil {
		err = fmt.Errorf("HMAC should have failed")
		FailWithError(t, err)
	}

	if bytes.Equal(decrypted, m) {
		err = fmt.Errorf("decryption should not have succeeded")
		FailWithError(t, err)
	}
}

// Test Zeroise, which is used in the EncryptReader
func TestZeroise(t *testing.T) {
	var err error
	var testVector = []byte("hello, world")

	if len(testVector) != len("hello, world") {
		err = fmt.Errorf("testVector improperly initialised")
		FailWithError(t, err)
	}

	Zeroise(&testVector)
	if len(testVector) != 0 {
		err = fmt.Errorf("testVector not empty after Zeroise")
		FailWithError(t, err)
	}
}

// Benchmark the generation of session keys.
func BenchmarkGenerateKey(b *testing.B) {
	for i := 0; i < b.N; i++ {
		key, err := GenerateKey()
		if err != nil || len(key) != KeySize {
			b.FailNow()
		}
		Zeroise(&key)
	}
}

// Benchmark the generation of long-term encryption keys.
func BenchmarkGenerateLTKey(b *testing.B) {
	for i := 0; i < b.N; i++ {
		key, err := GenerateLTKey()
		if err != nil || len(key) != KeySize {
			b.FailNow()
		}
		Zeroise(&key)
	}
}

// Benchmark the generation of initialisation vectors.
func BenchmarkGenerateIV(b *testing.B) {
	for i := 0; i < b.N; i++ {
		iv, err := GenerateIV()
		if err != nil || len(iv) != BlockSize {
			b.FailNow()
		}
		Zeroise(&iv)
	}
}

// Benchmark the encryption and decryption of a single block.
func BenchmarkEncryptBlock(b *testing.B) {
	for i := 0; i < b.N; i++ {
		m := []byte("Hello, world.")

		key, err := GenerateKey()
		if err != nil {
			fmt.Println("symkey")
			fmt.Println(err.Error())
			b.FailNow()
		}

		e, err := Encrypt(key, m)
		if err != nil {
			fmt.Println("encrypt")
			b.FailNow()
		}

		decrypted, err := Decrypt(key, e)
		if err != nil {
			fmt.Println("decrypt")
			b.FailNow()
		}

		if !bytes.Equal(decrypted, m) {
			fmt.Println("equal")
			b.FailNow()
		}

		Zeroise(&key)
	}
}

// Benchmark the scrubbing of a key.
func BenchmarkScrubKey(b *testing.B) {
	for i := 0; i < b.N; i++ {
		key, err := GenerateKey()
		if err != nil {
			b.FailNow()
		}

		err = Scrub(key, 3)
		if err != nil {
			b.FailNow()
		}
	}
}

// Benchmark encrypting and decrypting to bytes.
func BenchmarkByteCrypt(b *testing.B) {
	msg := []byte("Hello, world. Hallo, welt. Hej, världen.")

	for i := 0; i < b.N; i++ {
		enc, err := Encrypt(testKey, msg)
		if err != nil {
			b.FailNow()
		}

		dec, err := Decrypt(testKey, enc)
		if err != nil {
			b.FailNow()
		}

		if !bytes.Equal(msg, dec) {
			b.FailNow()
		}
	}
}
<commit_msg>clean up tests<commit_after>package authsym

import (
	"bytes"
	"fmt"
	"testing"
)

const (
	testEnc = "/tmp/test.out"
	testOut = "/tmp/test.dat"
	testRef = "testdata/vector01.dat"
	rPadRef = "testdata/vector02.dat"
)

var (
	testKey []byte
)

// FailWithError is a utility for dumping errors and failing the test.
func FailWithError(t *testing.T, err error) {
	fmt.Println("failed")
	if err != nil {
		fmt.Println("[!] ", err.Error())
	}
	t.FailNow()
}

// Test session key generation.
func TestGenerateKey(t *testing.T) {
	key, err := GenerateKey()
	if err != nil || len(key) != KeySize {
		FailWithError(t, err)
	}
	testKey = key
}

// Test long term key generation.
func TestGenerateLTKey(t *testing.T) {
	key, err := GenerateLTKey()
	if err != nil || len(key) != KeySize {
		FailWithError(t, err)
	}
}

// Test initialisation vector generation.
func TestGenerateIV(t *testing.T) {
	iv, err := GenerateIV()
	if err != nil || len(iv) != BlockSize {
		FailWithError(t, err)
	}
}

// Does D(E(k,m)) == m?
func TestEncryptDecryptBlock(t *testing.T) {
	m := []byte("Hello, world.")
	key, err := GenerateKey()
	if err != nil {
		FailWithError(t, err)
	}

	e, err := Encrypt(key, m)
	if err != nil {
		FailWithError(t, err)
	}

	decrypted, err := Decrypt(key, e)
	if err != nil {
		FailWithError(t, err)
	}

	if !bytes.Equal(decrypted, m) {
		err = fmt.Errorf("plaintext doesn't match original message")
		FailWithError(t, err)
	}
}

func TestEncryptDecryptBlockFails(t *testing.T) {
	m := []byte("Hello, world.")

	key, err := GenerateKey()
	if err != nil {
		FailWithError(t, err)
	}

	e, err := Encrypt(key, m)
	if err != nil {
		FailWithError(t, err)
	}

	n := len(e) - 2
	orig := e[n]
	if e[n] == 255 {
		e[n] = 0
	} else {
		e[n]++
	}
	if e[n] == orig {
		err = fmt.Errorf("byte not modified")
		FailWithError(t, err)
	}
	decrypted, err := Decrypt(key, e)
	if err == nil {
		err = fmt.Errorf("HMAC should have failed")
		FailWithError(t, err)
	}

	if bytes.Equal(decrypted, m) {
		err = fmt.Errorf("decryption should not have succeeded")
		FailWithError(t, err)
	}
}

// Test Zeroise, which is used in the EncryptReader
func TestZeroise(t *testing.T) {
	var err error
	var testVector = []byte("hello, world")

	if len(testVector) != len("hello, world") {
		err = fmt.Errorf("testVector improperly initialised")
		FailWithError(t, err)
	}

	Zeroise(&testVector)
	if len(testVector) != 0 {
		err = fmt.Errorf("testVector not empty after Zeroise")
		FailWithError(t, err)
	}
}

// Benchmark the generation of session keys.
func BenchmarkGenerateKey(b *testing.B) {
	for i := 0; i < b.N; i++ {
		key, err := GenerateKey()
		if err != nil || len(key) != KeySize {
			b.FailNow()
		}
		Zeroise(&key)
	}
}

// Benchmark the generation of long-term encryption keys.
func BenchmarkGenerateLTKey(b *testing.B) {
	for i := 0; i < b.N; i++ {
		key, err := GenerateLTKey()
		if err != nil || len(key) != KeySize {
			b.FailNow()
		}
		Zeroise(&key)
	}
}

// Benchmark the generation of initialisation vectors.
func BenchmarkGenerateIV(b *testing.B) {
	for i := 0; i < b.N; i++ {
		iv, err := GenerateIV()
		if err != nil || len(iv) != BlockSize {
			b.FailNow()
		}
		Zeroise(&iv)
	}
}

// Benchmark the encryption and decryption of a single block.
func BenchmarkEncryptBlock(b *testing.B) {
	for i := 0; i < b.N; i++ {
		m := []byte("Hello, world.")

		key, err := GenerateKey()
		if err != nil {
			b.FailNow()
		}

		e, err := Encrypt(key, m)
		if err != nil {
			b.FailNow()
		}

		decrypted, err := Decrypt(key, e)
		if err != nil {
			b.FailNow()
		}

		if !bytes.Equal(decrypted, m) {
			b.FailNow()
		}

		Zeroise(&key)
	}
}

// Benchmark the scrubbing of a key.
func BenchmarkScrubKey(b *testing.B) {
	for i := 0; i < b.N; i++ {
		key, err := GenerateKey()
		if err != nil {
			b.FailNow()
		}

		err = Scrub(key, 3)
		if err != nil {
			b.FailNow()
		}
	}
}

// Benchmark encrypting and decrypting to bytes.
func BenchmarkByteCrypt(b *testing.B) {
	msg := []byte("Hello, world. Hallo, welt. Hej, världen.")

	for i := 0; i < b.N; i++ {
		enc, err := Encrypt(testKey, msg)
		if err != nil {
			b.FailNow()
		}

		dec, err := Decrypt(testKey, enc)
		if err != nil {
			b.FailNow()
		}

		if !bytes.Equal(msg, dec) {
			b.FailNow()
		}
	}
}

// Benchmark encrypting and decrypting to bytes.
func BenchmarkByteUnauthCrypt(b *testing.B) {
	msg := []byte("Hello, world. Hallo, welt. Hej, världen.")

	for i := 0; i < b.N; i++ {
		enc, err := encrypt(testKey, msg)
		if err != nil {
			b.FailNow()
		}

		dec, err := decrypt(testKey, enc)
		if err != nil {
			b.FailNow()
		}

		if !bytes.Equal(msg, dec) {
			b.FailNow()
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	flag "github.com/ogier/pflag"
	"io/ioutil"
	"os"
	"strings"
	"text/template"
)

const helpStr = `    jgtr - JSON Go Template Renderer

USAGE:
    jgtr [OPTIONS]

    jgtr consumes a data file and a template file written in Go's text/template
    language. The values available in the data file are then used to render the
    template file and generate output.

    By default, jgtr reads the data and template from stdin, and writes output
    to stdout. Note that data and template cannot both come from stdin - at
    least one of the two must be specified via an option.

    jgtr can consume data from JSON, YAML 1.1 or TOML v0.2.0. You can specify
    which type to use via an option. If no such option is given, jgtr attempts
    to guess from the extension of the data file (if any). If the format is
    still ambiguous, jgtr uses JSON as the default.

OPTIONS:
    -d FILE, --data=FILE
        Read data data from FILE. Specify "-" (the default) to use stdin.

    -t FILE, --template=FILE
        Read template from FILE. Specify "-" (the default) to use stdin.

    -o FILE, --output=FILE
        Write rendered template to FILE. Specify "-" (the default) to use
        stdout.

    -j, --json
    	Specify the data format as JSON (default).

    -y, --yaml
    	Specify the data format as YAML.

    -T, --toml
    	Specify the data format as TOML.

    -h, --help
        Display this help.

    -V, --version
        Display jgtr version.`

const versionStr = `0.7.0`

func main() {
	help := flag.BoolP("help", "h", false, "show help")
	version := flag.BoolP("version", "V", false, "show version")

	dataPath := flag.StringP("data", "d", "-", "data file (JSON by default)")
	tmplPath := flag.StringP("template", "t", "-", "Go template file")
	outPath := flag.StringP("output", "o", "-", "output file")

	jsonFlag := flag.BoolP("json", "j", false, "interpret data as JSON")
	yamlFlag := flag.BoolP("yaml", "y", false, "interpret data as YAML")
	tomlFlag := flag.BoolP("toml", "T", false, "interpret data as TOML")

	flag.Parse()

	if *help {
		fmt.Println(helpStr)
		return
	}
	if *version {
		println(versionStr)
		return
	}

	if *dataPath == "-" && *tmplPath == "-" {
		println("Cannot use stdin for data and template simultaneously")
		os.Exit(1)
	}

	var data interface{} = nil
	var err error = nil
	if *yamlFlag { // check the flags first
		data, err = loadYAMLData(*dataPath)
	} else if *tomlFlag {
		data, err = loadTOMLData(*dataPath)
	} else if *jsonFlag {
		data, err = loadJSONData(*dataPath)
	} else if strings.HasSuffix(*dataPath, ".yaml") || strings.HasSuffix(*dataPath, ".yml") { // no flag? check the extension
		data, err = loadYAMLData(*dataPath)
	} else if strings.HasSuffix(*dataPath, ".toml") {
		data, err = loadTOMLData(*dataPath)
	} else { // default case: json
		data, err = loadJSONData(*dataPath)
	}
	if err != nil {
		panic(err)
	}

	tmpl, err := loadGoTemplate(*tmplPath)
	if err != nil {
		panic(err)
	}

	outFile, err := createStream(*outPath)
	if err != nil {
		panic(err)
	}
	defer closeStream(outFile)

	err = tmpl.Execute(outFile, data)
	if err != nil {
		panic(err)
	}
}

// loadGoTemplate parses a Go text template from the file specified by path.
// The file contents are parsed into a top-level template with the name "root".
// If the path is "-", then the template will be parsed from os.Stdin.
func loadGoTemplate(path string) (tmpl *template.Template, err error) {
	file, err := openStream(path)
	if err != nil {
		return
	}
	defer closeStream(file)

	rawTmpl, err := ioutil.ReadAll(file)
	if err != nil {
		return
	}

	// explicitly parse the raw string rather than using ParseFiles
	// ParseFiles creates templates whose names are those of the files
	// and associates them with the parent template
	// this creates some confusing behavior in Template.Parse
	// see http://stackoverflow.com/questions/11805356/text-template-issue-parse-vs-parsefiles
	// also note: functions have to be added before the template is parsed
	tmpl, err = template.New("root").Funcs(tmplFuncs).Parse(string(rawTmpl))
	return // again, whether err==nil or not, this is finished
}

// openStream behaves like os.Open, except that if the path is "-", then it
// simply returns os.Stdin.
func openStream(path string) (file *os.File, err error) {
	if path == "-" {
		file = os.Stdin
	} else {
		file, err = os.Open(path)
	}
	return
}

// createStream behaves like os.Create, except that if the path is "-", then it
// simply returns os.Stdout.
func createStream(path string) (file *os.File, err error) {
	if path == "-" {
		file = os.Stdout
	} else {
		file, err = os.Create(path)
	}
	return
}

// closeStream behaves like file.Close, except that if the file is os.Stdin or
// os.Stdout, it does nothing.
func closeStream(file *os.File) (err error) {
	if file == os.Stdout || file == os.Stdin {
		return
	}
	return file.Close()
}
<commit_msg>Separate extension detection from format flag checking<commit_after>package main

import (
	"fmt"
	flag "github.com/ogier/pflag"
	"io/ioutil"
	"os"
	"strings"
	"text/template"
)

const helpStr = `    jgtr - JSON Go Template Renderer

USAGE:
    jgtr [OPTIONS]

    jgtr consumes a data file and a template file written in Go's text/template
    language. The values available in the data file are then used to render the
    template file and generate output.

    By default, jgtr reads the data and template from stdin, and writes output
    to stdout. Note that data and template cannot both come from stdin - at
    least one of the two must be specified via an option.

    jgtr can consume data from JSON, YAML 1.1 or TOML v0.2.0. You can specify
    which type to use via an option. If no such option is given, jgtr attempts
    to guess from the extension of the data file (if any). If the format is
    still ambiguous, jgtr uses JSON as the default.

OPTIONS:
    -d FILE, --data=FILE
        Read data data from FILE. Specify "-" (the default) to use stdin.

    -t FILE, --template=FILE
        Read template from FILE. Specify "-" (the default) to use stdin.

    -o FILE, --output=FILE
        Write rendered template to FILE. Specify "-" (the default) to use
        stdout.

    -j, --json
    	Specify the data format as JSON (default).

    -y, --yaml
    	Specify the data format as YAML.

    -T, --toml
    	Specify the data format as TOML.

    -h, --help
        Display this help.

    -V, --version
        Display jgtr version.`

const versionStr = `0.7.0`

func main() {
	help := flag.BoolP("help", "h", false, "show help")
	version := flag.BoolP("version", "V", false, "show version")

	dataPath := flag.StringP("data", "d", "-", "data file (JSON by default)")
	tmplPath := flag.StringP("template", "t", "-", "Go template file")
	outPath := flag.StringP("output", "o", "-", "output file")

	jsonFlag := flag.BoolP("json", "j", false, "interpret data as JSON")
	yamlFlag := flag.BoolP("yaml", "y", false, "interpret data as YAML")
	tomlFlag := flag.BoolP("toml", "T", false, "interpret data as TOML")

	flag.Parse()

	if *help {
		fmt.Println(helpStr)
		return
	}
	if *version {
		println(versionStr)
		return
	}

	if *dataPath == "-" && *tmplPath == "-" {
		println("Cannot use stdin for data and template simultaneously")
		os.Exit(1)
	}

	// no flag set - check extensions
	if !(*jsonFlag || *yamlFlag || *tomlFlag) {
		if strings.HasSuffix(*dataPath, ".yaml") || strings.HasSuffix(*dataPath, ".yml") {
			flag.Set("yaml", "true")
		} else if strings.HasSuffix(*dataPath, ".toml") {
			flag.Set("toml", "true")
		} else {
			flag.Set("json", "true")
		}
	}

	var data interface{} = nil
	var err error = nil
	if *yamlFlag {
		data, err = loadYAMLData(*dataPath)
	} else if *tomlFlag {
		data, err = loadTOMLData(*dataPath)
	} else {
		data, err = loadJSONData(*dataPath)
	}
	if err != nil {
		panic(err)
	}

	tmpl, err := loadGoTemplate(*tmplPath)
	if err != nil {
		panic(err)
	}

	outFile, err := createStream(*outPath)
	if err != nil {
		panic(err)
	}
	defer closeStream(outFile)

	err = tmpl.Execute(outFile, data)
	if err != nil {
		panic(err)
	}
}

// loadGoTemplate parses a Go text template from the file specified by path.
// The file contents are parsed into a top-level template with the name "root".
// If the path is "-", then the template will be parsed from os.Stdin.
func loadGoTemplate(path string) (tmpl *template.Template, err error) {
	file, err := openStream(path)
	if err != nil {
		return
	}
	defer closeStream(file)

	rawTmpl, err := ioutil.ReadAll(file)
	if err != nil {
		return
	}

	// explicitly parse the raw string rather than using ParseFiles
	// ParseFiles creates templates whose names are those of the files
	// and associates them with the parent template
	// this creates some confusing behavior in Template.Parse
	// see http://stackoverflow.com/questions/11805356/text-template-issue-parse-vs-parsefiles
	// also note: functions have to be added before the template is parsed
	tmpl, err = template.New("root").Funcs(tmplFuncs).Parse(string(rawTmpl))
	return // again, whether err==nil or not, this is finished
}

// openStream behaves like os.Open, except that if the path is "-", then it
// simply returns os.Stdin.
func openStream(path string) (file *os.File, err error) {
	if path == "-" {
		file = os.Stdin
	} else {
		file, err = os.Open(path)
	}
	return
}

// createStream behaves like os.Create, except that if the path is "-", then it
// simply returns os.Stdout.
func createStream(path string) (file *os.File, err error) {
	if path == "-" {
		file = os.Stdout
	} else {
		file, err = os.Create(path)
	}
	return
}

// closeStream behaves like file.Close, except that if the file is os.Stdin or
// os.Stdout, it does nothing.
func closeStream(file *os.File) (err error) {
	if file == os.Stdout || file == os.Stdin {
		return
	}
	return file.Close()
}
<|endoftext|>