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

package libkb

import (
	"fmt"
	"time"

	keybase1 "github.com/keybase/client/go/protocol/keybase1"
	jsonw "github.com/keybase/go-jsonw"
	"golang.org/x/net/context"
	"stathat.com/c/ramcache"
)

type ResolveResult struct {
	uid                keybase1.UID
	body               *jsonw.Wrapper
	err                error
	queriedKbUsername  string
	queriedByUID       bool
	resolvedKbUsername string
	cachedAt           time.Time
	mutable            bool
}

const (
	resolveCacheTTL           = 12 * time.Hour
	ResolveCacheMaxAge        = 12 * time.Hour
	ResolveCacheMaxAgeMutable = 20 * time.Minute
	resolveCacheMaxAgeErrored = 5 * time.Second
)

func (res *ResolveResult) GetUID() keybase1.UID {
	return res.uid
}

func (res *ResolveResult) GetUsername() string {
	return res.resolvedKbUsername
}
func (res *ResolveResult) GetNormalizedUsername() NormalizedUsername {
	return NewNormalizedUsername(res.GetUsername())
}
func (res *ResolveResult) GetNormalizedQueriedUsername() NormalizedUsername {
	return NewNormalizedUsername(res.queriedKbUsername)
}

func (res *ResolveResult) WasKBAssertion() bool {
	return res.queriedKbUsername != "" || res.queriedByUID
}

func (res *ResolveResult) GetError() error {
	return res.err
}

func (res *ResolveResult) GetBody() *jsonw.Wrapper {
	return res.body
}

func (r *Resolver) ResolveWithBody(input string) ResolveResult {
	return r.resolve(input, true)
}

func (r *Resolver) Resolve(input string) ResolveResult {
	return r.resolve(input, false)
}

func (r *Resolver) resolve(input string, withBody bool) (res ResolveResult) {
	defer r.G().Trace(fmt.Sprintf("Resolving username %q", input), func() error { return res.err })()

	var au AssertionURL
	if au, res.err = ParseAssertionURL(r.G().MakeAssertionContext(), input, false); res.err != nil {
		return res
	}
	res = r.resolveURL(context.TODO(), au, input, withBody, false)
	return res
}

func (r *Resolver) ResolveFullExpression(ctx context.Context, input string) (res ResolveResult) {
	return r.resolveFullExpression(ctx, input, false, false)
}

func (r *Resolver) ResolveFullExpressionNeedUsername(ctx context.Context, input string) (res ResolveResult) {
	return r.resolveFullExpression(ctx, input, false, true)
}

func (r *Resolver) ResolveFullExpressionWithBody(ctx context.Context, input string) (res ResolveResult) {
	return r.resolveFullExpression(ctx, input, true, false)
}

func (r *Resolver) resolveFullExpression(ctx context.Context, input string, withBody bool, needUsername bool) (res ResolveResult) {
	defer r.G().CTrace(ctx, fmt.Sprintf("Resolver#resolveFullExpression(%q)", input), func() error { return res.err })()

	var expr AssertionExpression
	expr, res.err = AssertionParseAndOnly(r.G().MakeAssertionContext(), input)
	if res.err != nil {
		return res
	}
	u := FindBestIdentifyComponentURL(expr)
	if u == nil {
		res.err = ResolutionError{Input: input, Msg: "Cannot find a resolvable factor"}
		return res
	}
	return r.resolveURL(ctx, u, input, withBody, needUsername)
}

func (r *Resolver) getFromDiskCache(ctx context.Context, key string) (ret *ResolveResult) {
	defer r.G().CTraceOK(ctx, fmt.Sprintf("Resolver#getFromDiskCache(%q)", key), func() bool { return ret != nil })()
	var uid keybase1.UID
	found, err := r.G().LocalDb.GetInto(&uid, resolveDbKey(key))
	r.Stats.diskGets++
	if err != nil {
		r.G().Log.CWarningf(ctx, "Problem fetching resolve result from local DB: %s", err)
		return nil
	}
	if !found {
		r.Stats.diskGetMisses++
		return nil
	}
	r.Stats.diskGetHits++
	return &ResolveResult{uid: uid}
}

func (r *Resolver) resolveURL(ctx context.Context, au AssertionURL, input string, withBody bool, needUsername bool) ResolveResult {

	// A standard keybase UID, so it's already resolved... unless we explicitly
	// need it!
	if !needUsername {
		if tmp := au.ToUID(); tmp.Exists() {
			return ResolveResult{uid: tmp}
		}
	}

	ck := au.CacheKey()

	if p := r.getFromMemCache(ctx, ck); p != nil && (!needUsername || len(p.resolvedKbUsername) > 0) {
		return *p
	}

	if p := r.getFromDiskCache(ctx, ck); p != nil && (!needUsername || len(p.resolvedKbUsername) > 0) {
		p.mutable = !au.IsKeybase()
		r.putToMemCache(ck, *p)
		return *p
	}

	res := r.resolveURLViaServerLookup(ctx, au, input, withBody)

	// Cache for a shorter period of time if it's not a Keybase identity
	res.mutable = !au.IsKeybase()
	r.putToMemCache(ck, res)
	r.putToDiskCache(ck, res)

	return res
}

func (r *Resolver) resolveURLViaServerLookup(ctx context.Context, au AssertionURL, input string, withBody bool) (res ResolveResult) {
	defer r.G().CTrace(ctx, fmt.Sprintf("Resolver#resolveURLViaServerLookup(input = %q)", input), func() error { return res.err })()

	var key, val string
	var ares *APIRes
	var l int

	if au.IsKeybase() {
		res.queriedKbUsername = au.GetValue()
	} else if au.IsUID() {
		res.queriedByUID = true
	}

	if key, val, res.err = au.ToLookup(); res.err != nil {
		return
	}

	ha := HTTPArgsFromKeyValuePair(key, S{val})
	ha.Add("multi", I{1})
	fields := "basics"
	if withBody {
		fields += ",public_keys,pictures"
	}
	ha.Add("fields", S{fields})
	ares, res.err = r.G().API.Get(APIArg{
		Endpoint:       "user/lookup",
		NeedSession:    false,
		Args:           ha,
		AppStatusCodes: []int{SCOk, SCNotFound, SCDeleted},
		NetContext:     ctx,
	})

	if res.err != nil {
		r.G().Log.CDebugf(ctx, "API user/lookup %q error: %s", input, res.err)
		return
	}
	switch ares.AppStatus.Code {
	case SCNotFound:
		r.G().Log.CDebugf(ctx, "API user/lookup %q not found", input)
		res.err = NotFoundError{}
		return
	case SCDeleted:
		r.G().Log.CDebugf(ctx, "API user/lookup %q deleted", input)
		res.err = DeletedError{Msg: fmt.Sprintf("user %q deleted", input)}
		return
	}

	var them *jsonw.Wrapper
	if them, res.err = ares.Body.AtKey("them").ToArray(); res.err != nil {
		return
	}

	if l, res.err = them.Len(); res.err != nil {
		return
	}

	if l == 0 {
		res.err = ResolutionError{Input: input, Msg: "No resolution found"}
	} else if l > 1 {
		res.err = ResolutionError{Input: input, Msg: "Identify is ambiguous"}
	} else {
		res.body = them.AtIndex(0)
		res.uid, res.err = GetUID(res.body.AtKey("id"))
		if res.err == nil {
			res.resolvedKbUsername, res.err = res.body.AtPath("basics.username").GetString()
		}
	}

	return
}

type ResolveCacheStats struct {
	misses          int
	timeouts        int
	mutableTimeouts int
	errorTimeouts   int
	hits            int
	diskGets        int
	diskGetHits     int
	diskGetMisses   int
	diskPuts        int
}

type Resolver struct {
	Contextified
	cache   *ramcache.Ramcache
	Stats   ResolveCacheStats
	NowFunc func() time.Time
}

func (s ResolveCacheStats) Eq(m, t, mt, et, h int) bool {
	return (s.misses == m) && (s.timeouts == t) && (s.mutableTimeouts == mt) && (s.errorTimeouts == et) && (s.hits == h)
}

func NewResolver(g *GlobalContext) *Resolver {
	return &Resolver{
		Contextified: NewContextified(g),
		cache:        nil,
		NowFunc:      func() time.Time { return time.Now() },
	}
}

func (r *Resolver) EnableCaching() {
	cache := ramcache.New()
	cache.MaxAge = ResolveCacheMaxAge
	cache.TTL = resolveCacheTTL
	r.cache = cache
}

func (r *Resolver) Shutdown() {
	if r.cache == nil {
		return
	}
	r.cache.Shutdown()
}

func (r *Resolver) getFromMemCache(ctx context.Context, key string) (ret *ResolveResult) {
	defer r.G().CTraceOK(ctx, fmt.Sprintf("Resolver#getFromMemCache(%q)", key), func() bool { return ret != nil })()
	if r.cache == nil {
		return nil
	}
	res, _ := r.cache.Get(key)
	if res == nil {
		r.Stats.misses++
		return nil
	}
	rres, ok := res.(*ResolveResult)
	if !ok {
		r.Stats.misses++
		return nil
	}
	now := r.NowFunc()
	if now.Sub(rres.cachedAt) > ResolveCacheMaxAge {
		r.Stats.timeouts++
		return nil
	}
	if rres.mutable && now.Sub(rres.cachedAt) > ResolveCacheMaxAgeMutable {
		r.Stats.mutableTimeouts++
		return nil
	}
	if rres.err != nil && now.Sub(rres.cachedAt) > resolveCacheMaxAgeErrored {
		r.Stats.errorTimeouts++
		return nil
	}
	r.Stats.hits++
	return rres
}

func resolveDbKey(key string) DbKey {
	return DbKey{
		Typ: DBResolveUsernameToUID,
		Key: NewNormalizedUsername(key).String(),
	}
}

func (r *Resolver) putToDiskCache(key string, res ResolveResult) {
	// Only cache immutable resolutions to disk
	if res.mutable {
		return
	}
	r.Stats.diskPuts++
	err := r.G().LocalDb.PutObj(resolveDbKey(key), nil, res.uid)
	if err != nil {
		r.G().Log.Warning("Cannot put resolve result to disk: %s", err)
	}
}

// Put receives a copy of a ResolveResult, clears out the body
// to avoid caching data that can go stale, and stores the result.
func (r *Resolver) putToMemCache(key string, res ResolveResult) {
	if r.cache == nil {
		return
	}
	res.cachedAt = r.NowFunc()
	res.body = nil // Don't cache body
	r.cache.Set(key, &res)
}
<commit_msg>do not cache errors in UID resolver<commit_after>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

package libkb

import (
	"fmt"
	"time"

	keybase1 "github.com/keybase/client/go/protocol/keybase1"
	jsonw "github.com/keybase/go-jsonw"
	"golang.org/x/net/context"
	"stathat.com/c/ramcache"
)

type ResolveResult struct {
	uid                keybase1.UID
	body               *jsonw.Wrapper
	err                error
	queriedKbUsername  string
	queriedByUID       bool
	resolvedKbUsername string
	cachedAt           time.Time
	mutable            bool
}

const (
	resolveCacheTTL           = 12 * time.Hour
	ResolveCacheMaxAge        = 12 * time.Hour
	ResolveCacheMaxAgeMutable = 20 * time.Minute
	resolveCacheMaxAgeErrored = 5 * time.Second
)

func (res *ResolveResult) GetUID() keybase1.UID {
	return res.uid
}

func (res *ResolveResult) GetUsername() string {
	return res.resolvedKbUsername
}
func (res *ResolveResult) GetNormalizedUsername() NormalizedUsername {
	return NewNormalizedUsername(res.GetUsername())
}
func (res *ResolveResult) GetNormalizedQueriedUsername() NormalizedUsername {
	return NewNormalizedUsername(res.queriedKbUsername)
}

func (res *ResolveResult) WasKBAssertion() bool {
	return res.queriedKbUsername != "" || res.queriedByUID
}

func (res *ResolveResult) GetError() error {
	return res.err
}

func (res *ResolveResult) GetBody() *jsonw.Wrapper {
	return res.body
}

func (r *Resolver) ResolveWithBody(input string) ResolveResult {
	return r.resolve(input, true)
}

func (r *Resolver) Resolve(input string) ResolveResult {
	return r.resolve(input, false)
}

func (r *Resolver) resolve(input string, withBody bool) (res ResolveResult) {
	defer r.G().Trace(fmt.Sprintf("Resolving username %q", input), func() error { return res.err })()

	var au AssertionURL
	if au, res.err = ParseAssertionURL(r.G().MakeAssertionContext(), input, false); res.err != nil {
		return res
	}
	res = r.resolveURL(context.TODO(), au, input, withBody, false)
	return res
}

func (r *Resolver) ResolveFullExpression(ctx context.Context, input string) (res ResolveResult) {
	return r.resolveFullExpression(ctx, input, false, false)
}

func (r *Resolver) ResolveFullExpressionNeedUsername(ctx context.Context, input string) (res ResolveResult) {
	return r.resolveFullExpression(ctx, input, false, true)
}

func (r *Resolver) ResolveFullExpressionWithBody(ctx context.Context, input string) (res ResolveResult) {
	return r.resolveFullExpression(ctx, input, true, false)
}

func (r *Resolver) resolveFullExpression(ctx context.Context, input string, withBody bool, needUsername bool) (res ResolveResult) {
	defer r.G().CTrace(ctx, fmt.Sprintf("Resolver#resolveFullExpression(%q)", input), func() error { return res.err })()

	var expr AssertionExpression
	expr, res.err = AssertionParseAndOnly(r.G().MakeAssertionContext(), input)
	if res.err != nil {
		return res
	}
	u := FindBestIdentifyComponentURL(expr)
	if u == nil {
		res.err = ResolutionError{Input: input, Msg: "Cannot find a resolvable factor"}
		return res
	}
	return r.resolveURL(ctx, u, input, withBody, needUsername)
}

func (r *Resolver) getFromDiskCache(ctx context.Context, key string) (ret *ResolveResult) {
	defer r.G().CTraceOK(ctx, fmt.Sprintf("Resolver#getFromDiskCache(%q)", key), func() bool { return ret != nil })()
	var uid keybase1.UID
	found, err := r.G().LocalDb.GetInto(&uid, resolveDbKey(key))
	r.Stats.diskGets++
	if err != nil {
		r.G().Log.CWarningf(ctx, "Problem fetching resolve result from local DB: %s", err)
		return nil
	}
	if !found {
		r.Stats.diskGetMisses++
		return nil
	}
	r.Stats.diskGetHits++
	return &ResolveResult{uid: uid}
}

func (r *Resolver) resolveURL(ctx context.Context, au AssertionURL, input string, withBody bool, needUsername bool) ResolveResult {

	// A standard keybase UID, so it's already resolved... unless we explicitly
	// need it!
	if !needUsername {
		if tmp := au.ToUID(); tmp.Exists() {
			return ResolveResult{uid: tmp}
		}
	}

	ck := au.CacheKey()

	if p := r.getFromMemCache(ctx, ck); p != nil && (!needUsername || len(p.resolvedKbUsername) > 0) {
		return *p
	}

	if p := r.getFromDiskCache(ctx, ck); p != nil && (!needUsername || len(p.resolvedKbUsername) > 0) {
		p.mutable = !au.IsKeybase()
		r.putToMemCache(ck, *p)
		return *p
	}

	res := r.resolveURLViaServerLookup(ctx, au, input, withBody)

	// Cache for a shorter period of time if it's not a Keybase identity
	res.mutable = !au.IsKeybase()
	r.putToMemCache(ck, res)
	r.putToDiskCache(ck, res)

	return res
}

func (r *Resolver) resolveURLViaServerLookup(ctx context.Context, au AssertionURL, input string, withBody bool) (res ResolveResult) {
	defer r.G().CTrace(ctx, fmt.Sprintf("Resolver#resolveURLViaServerLookup(input = %q)", input), func() error { return res.err })()

	var key, val string
	var ares *APIRes
	var l int

	if au.IsKeybase() {
		res.queriedKbUsername = au.GetValue()
	} else if au.IsUID() {
		res.queriedByUID = true
	}

	if key, val, res.err = au.ToLookup(); res.err != nil {
		return
	}

	ha := HTTPArgsFromKeyValuePair(key, S{val})
	ha.Add("multi", I{1})
	fields := "basics"
	if withBody {
		fields += ",public_keys,pictures"
	}
	ha.Add("fields", S{fields})
	ares, res.err = r.G().API.Get(APIArg{
		Endpoint:       "user/lookup",
		NeedSession:    false,
		Args:           ha,
		AppStatusCodes: []int{SCOk, SCNotFound, SCDeleted},
		NetContext:     ctx,
	})

	if res.err != nil {
		r.G().Log.CDebugf(ctx, "API user/lookup %q error: %s", input, res.err)
		return
	}
	switch ares.AppStatus.Code {
	case SCNotFound:
		r.G().Log.CDebugf(ctx, "API user/lookup %q not found", input)
		res.err = NotFoundError{}
		return
	case SCDeleted:
		r.G().Log.CDebugf(ctx, "API user/lookup %q deleted", input)
		res.err = DeletedError{Msg: fmt.Sprintf("user %q deleted", input)}
		return
	}

	var them *jsonw.Wrapper
	if them, res.err = ares.Body.AtKey("them").ToArray(); res.err != nil {
		return
	}

	if l, res.err = them.Len(); res.err != nil {
		return
	}

	if l == 0 {
		res.err = ResolutionError{Input: input, Msg: "No resolution found"}
	} else if l > 1 {
		res.err = ResolutionError{Input: input, Msg: "Identify is ambiguous"}
	} else {
		res.body = them.AtIndex(0)
		res.uid, res.err = GetUID(res.body.AtKey("id"))
		if res.err == nil {
			res.resolvedKbUsername, res.err = res.body.AtPath("basics.username").GetString()
		}
	}

	return
}

type ResolveCacheStats struct {
	misses          int
	timeouts        int
	mutableTimeouts int
	errorTimeouts   int
	hits            int
	diskGets        int
	diskGetHits     int
	diskGetMisses   int
	diskPuts        int
}

type Resolver struct {
	Contextified
	cache   *ramcache.Ramcache
	Stats   ResolveCacheStats
	NowFunc func() time.Time
}

func (s ResolveCacheStats) Eq(m, t, mt, et, h int) bool {
	return (s.misses == m) && (s.timeouts == t) && (s.mutableTimeouts == mt) && (s.errorTimeouts == et) && (s.hits == h)
}

func NewResolver(g *GlobalContext) *Resolver {
	return &Resolver{
		Contextified: NewContextified(g),
		cache:        nil,
		NowFunc:      func() time.Time { return time.Now() },
	}
}

func (r *Resolver) EnableCaching() {
	cache := ramcache.New()
	cache.MaxAge = ResolveCacheMaxAge
	cache.TTL = resolveCacheTTL
	r.cache = cache
}

func (r *Resolver) Shutdown() {
	if r.cache == nil {
		return
	}
	r.cache.Shutdown()
}

func (r *Resolver) getFromMemCache(ctx context.Context, key string) (ret *ResolveResult) {
	defer r.G().CTraceOK(ctx, fmt.Sprintf("Resolver#getFromMemCache(%q)", key), func() bool { return ret != nil })()
	if r.cache == nil {
		return nil
	}
	res, _ := r.cache.Get(key)
	if res == nil {
		r.Stats.misses++
		return nil
	}
	rres, ok := res.(*ResolveResult)
	if !ok {
		r.Stats.misses++
		return nil
	}
	now := r.NowFunc()
	if now.Sub(rres.cachedAt) > ResolveCacheMaxAge {
		r.Stats.timeouts++
		return nil
	}
	if rres.mutable && now.Sub(rres.cachedAt) > ResolveCacheMaxAgeMutable {
		r.Stats.mutableTimeouts++
		return nil
	}
	if rres.err != nil && now.Sub(rres.cachedAt) > resolveCacheMaxAgeErrored {
		r.Stats.errorTimeouts++
		return nil
	}
	r.Stats.hits++
	return rres
}

func resolveDbKey(key string) DbKey {
	return DbKey{
		Typ: DBResolveUsernameToUID,
		Key: NewNormalizedUsername(key).String(),
	}
}

func (r *Resolver) putToDiskCache(key string, res ResolveResult) {
	// Only cache immutable resolutions to disk
	if res.mutable || res.err != nil {
		return
	}
	r.Stats.diskPuts++
	err := r.G().LocalDb.PutObj(resolveDbKey(key), nil, res.uid)
	if err != nil {
		r.G().Log.Warning("Cannot put resolve result to disk: %s", err)
	}
}

// Put receives a copy of a ResolveResult, clears out the body
// to avoid caching data that can go stale, and stores the result.
func (r *Resolver) putToMemCache(key string, res ResolveResult) {
	if r.cache == nil {
		return
	}
	if res.err == nil {
		return
	}
	res.cachedAt = r.NowFunc()
	res.body = nil // Don't cache body
	r.cache.Set(key, &res)
}
<|endoftext|>


<commit_before>package tricorder

import (
	"errors"
	"github.com/Symantec/tricorder/go/tricorder/units"
)

var (
	// RegisterMetric returns this if given path is already in use.
	ErrPathInUse = errors.New("tricorder: Path in use")
)

// A region represents a collection of variables for metrics that are all
// updated by a common function. Each time a client sends a request for one or
// more metrics backed by variables within a particular region, tricorder
// calls that region’s update function one time before reading any of the
// variables in that region to to respond to the client. However, to provide
// a consistent view of the variables within a region, tricorder will never
// call a region’s update function once it has begun reading variables in that
// region to service an in-process request.  If tricorder does happen to
// receive an incoming request for metrics from a given region after tricorder
// has begun reading variables in that same region to service another
// in-process request, tricorder will skip calling the region’s update
// function for the incoming request. In this case, the two requests will
// read the same data from that region.
type Region region

// NewRegion creates a new region with a particular update function
func RegisterRegion(updateFunc func()) *Region {
	return (*Region)(newRegion(updateFunc))
}

// RegisterMetric registers a single metric with the health system.
// path is the absolute path of the metric e.g "/proc/rpc";
// metric is the metric to register;
// unit is the unit of measurement for the metric;
// description is the description of the metric.
// RegisterMetric returns ErrPathInUse if path already represents a metric
// or a directory.
// RegisterMetric panics if metric is not of a valid type.
func RegisterMetric(
	path string,
	metric interface{},
	unit units.Unit,
	description string) error {
	return root.registerMetric(
		newPathSpec(path), metric, nil, unit, description)
}

// RegisterMetricWithRegion works just like RegisterMetrics but allows
// the caller to specify the region to which the variable or callback function
// being registered belongs. RegisterMetricWithRegion ignores the region
// parameter when registering a distribution.
func RegisterMetricInRegion(
	path string,
	metric interface{},
	r *Region,
	unit units.Unit,
	description string) error {
	return root.registerMetric(newPathSpec(path), metric, (*region)(r), unit, description)
}

// UnregisterPath unregisters the metric or DirectorySpec at the given path.
// UnregisterPath ignores requests to unregister the root path.
func UnregisterPath(path string) {
	root.unregisterPath(newPathSpec(path))
}

// Bucketer represents the organization of values into buckets for
// distributions. Bucketer instances are immutable.
type Bucketer struct {
	pieces []*bucketPiece
}

var (
	// Ranges in powers of two
	PowersOfTwo = NewExponentialBucketer(20, 1.0, 2.0)
	// Ranges in powers of four
	PowersOfFour = NewExponentialBucketer(11, 1.0, 4.0)
	// Ranges in powers of 10
	PowersOfTen = NewExponentialBucketer(7, 1.0, 10.0)
)

// NewExponentialBucketer returns a Bucketer representing buckets on
// a geometric scale. NewExponentialBucketer(25, 3.0, 1.7) means 25 buckets
// starting with <3.0; 3.0 - 5.1; 5.1 - 8.67; 8.67 - 14.739 etc.
// NewExponentialBucketer panics if count < 2 or if start <= 0 or if scale <= 1.
func NewExponentialBucketer(count int, start, scale float64) *Bucketer {
	return newBucketerFromEndpoints(
		newExponentialBucketerStream(count, start, scale))
}

// NewLinearBucketer returns a Bucketer representing bucktes on
// a linear scale. NewLinearBucketer(5, 0, 10) means 5 buckets
// starting with <0; 0-10; 10-20; 20-30; >=30.
// NewLinearBucketer panics if count < 2 or if increment <= 0.
func NewLinearBucketer(count int, start, increment float64) *Bucketer {
	return newBucketerFromEndpoints(
		newLinearBucketerStream(count, start, increment))
}

// NewArbitraryBucketer returns a Bucketer representing specific endpoints
// NewArbitraryBucketer(10.0, 20.0, 30.0) means 4 buckets:
// <10.0; 10.0 - 20.0; 20.0 - 30.0; >= 30.0.
// NewArbitraryBucketer panics if it is called with no arguments.
// It is the caller's responsibility to ensure that the arguments are in
// ascending order.
func NewArbitraryBucketer(endpoints ...float64) *Bucketer {
	return newBucketerFromEndpoints(endpoints)
}

// NewGeometricBucketer returns a Bucketer representing endpoints
// of the form 10^k, 2*10^k, 5*10^k. lower is the lower bound of
// the endpoints; upper is the upper bound of the endpoints.
// NewGeometricBucker(0.5, 50) ==>
// <0.5; 0.5-1; 1-2; 2-5; 5-10; 10-20; 20-50; >=50
func NewGeometricBucketer(lower, upper float64) *Bucketer {
	return newBucketerFromEndpoints(
		newGeometricBucketerStream(lower, upper))
}

// NewCumulativeDistribution creates a new CumulativeDistribution that uses
// this bucketer to distribute values.
func (b *Bucketer) NewCumulativeDistribution() *CumulativeDistribution {
	return (*CumulativeDistribution)(newDistribution(b, false))
}

// NewNonCumulativeDistribution creates a new NonCumulativeDistribution that
// uses this bucketer to distribute values.
func (b *Bucketer) NewNonCumulativeDistribution() *NonCumulativeDistribution {
	return (*NonCumulativeDistribution)(newDistribution(b, true))
}

// CumulativeDistribution represents a metric that is a distribution of
// values. Cumulative distributions only receive new values.
type CumulativeDistribution distribution

// Add adds a single value to this CumulativeDistribution instance.
// value can be a float32, float64, or a time.Duration.
// If a time.Duration, Add converts it to the same unit of time specified in
// the RegisterMetric call made to register this instance.
func (c *CumulativeDistribution) Add(value interface{}) {
	(*distribution)(c).Add(value)
}

// Unlike in CumulativeDistributions,values in NonCumulativeDistributions
// can change shifting from bucket to bucket.
type NonCumulativeDistribution distribution

// Add adds a single value to this NonCumulativeDistribution instance.
// value can be a float32, float64, or a time.Duration.
// If a time.Duration, Add converts it to the same unit of time specified in
// the RegisterMetric call made to register this instance.
func (c *NonCumulativeDistribution) Add(value interface{}) {
	(*distribution)(c).Add(value)
}

// Update updates a value in this NonCumulativeDistribution instance.
// oldValue and newValue can be a float32, float64, or a time.Duration.
// If a time.Duration, Update converts them to the same unit of time specified
// in the RegisterMetric call made to register this instance.
// The reliability of Update() depends on the caller providing the correct
// old value of what is being changed. Failure to do this results in
// undefined behavior.
// Update updates all distribution statistics in the expected way; however,
// it updates min and max such that min only gets smaller and max only gets
// larger. If update Updates values such that they fall into a narrower
// range than before, min and max remain unchanged to indicate the all-time
// min and all-time max. To have min and max reflect the current min and max
// instead of the all-time min and max, see UpdateMinMax().
func (d *NonCumulativeDistribution) Update(oldValue, newValue interface{}) {
	(*distribution)(d).Update(oldValue, newValue)
}

// UpdateMinMax() estimates the current min and max of this distribution.
// and updates min and max accordingly.
// As this call may be expensive, clients need not use unless both are
// true:
// 1) the client has made calls to Update which narrowed the current
// min and max.
// 2) The client wants min and max to to reflect the current min and max
// instead of the all-time min and max.
// This method only estimates. The only guarantees that it makes
// upon returning are:
// original_min <= min <= current_min and old_max >= max >= current_max.
// In fact, calling this method may do nothing at all which would still be
// correct behavior.
func (d *NonCumulativeDistribution) UpdateMinMax() {
}

// Remove removes a value from this NonCumulativeDistribution instance.
// valueToBeRemoved can be a float32, float64, or a time.Duration.
// If a time.Duration, Remove converts it to the same unit of time
// specified in the RegisterMetric call made to register this instance.
// The reliability of Remove() depends on the caller providing a value
// already in the distribution. Failure to do this results in
// undefined behavior.
// Remove updates all distribution statistics in the expected way; however,
// it leaves min and max unchanged.
// To have min and max reflect the current min and max
// instead of the all-time min and max, see UpdateMinMax().
func (d *NonCumulativeDistribution) Remove(valueToBeRemoved interface{}) {
	(*distribution)(d).Remove(valueToBeRemoved)
}

// Sum returns the sum of the values in this distribution.
func (d *NonCumulativeDistribution) Sum() float64 {
	return (*distribution)(d).Sum()
}

// DirectorySpec represents a specific directory in the heirarchy of
// metrics.
type DirectorySpec directory

// RegisterDirectory returns the the DirectorySpec registered with path.
// If nothing is registered with path, RegisterDirectory registers a
// new DirectorySpec with path and returns it.
// RegisterDirectory returns ErrPathInUse if path is already associated
// with a metric.
func RegisterDirectory(path string) (dirSpec *DirectorySpec, err error) {
	r, e := root.registerDirectory(newPathSpec(path))
	return (*DirectorySpec)(r), e
}

// RegisterMetric works just like the package level RegisterMetric
// except that path is relative to this DirectorySpec.
func (d *DirectorySpec) RegisterMetric(
	path string,
	metric interface{},
	unit units.Unit,
	description string) error {
	return (*directory)(d).registerMetric(newPathSpec(path), metric, nil, unit, description)
}

// RegisterMetricWithRegion works just like the package level
// RegisterMetricWithRegion except that path is relative to this
// DirectorySpec.
func (d *DirectorySpec) RegisterMetricInRegion(
	path string,
	metric interface{},
	r *Region,
	unit units.Unit,
	description string) error {
	return (*directory)(d).registerMetric(newPathSpec(path), metric, (*region)(r), unit, description)
}

// RegisterDirectory works just like the package level RegisterDirectory
// except that path is relative to this DirectorySpec.
func (d *DirectorySpec) RegisterDirectory(
	path string) (dirSpec *DirectorySpec, err error) {
	r, e := (*directory)(d).registerDirectory(newPathSpec(path))
	return (*DirectorySpec)(r), e
}

// Returns the absolute path this object represents
func (d *DirectorySpec) AbsPath() string {
	return (*directory)(d).AbsPath()
}

// UnregisterPath works just like the package level UnregisterPath
// except that path is relative to this DirectorySpec.
func (d *DirectorySpec) UnregisterPath(path string) {
	(*directory)(d).unregisterPath(newPathSpec(path))
}

// UnregisterDirectory unregisters this DirectorySpec instance along with
// all metrics and directories within it. The caller can unregister any
// DirectorySpec instance except the one representing the top level directory.
// That DirectorySpec instance simply ignores calls to UnregisterDirectory.
// Using an unregistered DirectorySpec instance to register new metrics may
// cause a panic.
func (d *DirectorySpec) UnregisterDirectory() {
	(*directory)(d).unregisterDirectory()
}

// RegisterFlags registers each application flag as a metric under /proc/flags.
func RegisterFlags() {
	registerFlags()
}

// SetFlagUnit sets the unit for a specific flag. If the flag is a
// time.Duration, the default unit is units.Second; otherwise the default
// unit is units.None. If the client wishes to override the default unit for
// a flag, they call this after registering a flag but before calling
// ParseFlags.
func SetFlagUnit(name string, unit units.Unit) {
	setFlagUnit(name, unit)
}
<commit_msg>Minor comment fixes.<commit_after>package tricorder

import (
	"errors"
	"github.com/Symantec/tricorder/go/tricorder/units"
)

var (
	// RegisterMetric returns this if given path is already in use.
	ErrPathInUse = errors.New("tricorder: Path in use")
)

// A region represents a collection of variables for metrics that are all
// updated by a common function. Each time a client sends a request for one or
// more metrics backed by variables within a particular region, tricorder
// calls that region’s update function one time before reading any of the
// variables in that region to respond to the client. However, to provide
// a consistent view of the variables within a region, tricorder will never
// call a region’s update function once it has begun reading variables in that
// region to service an in-process request.  If tricorder does happen to
// receive an incoming request for metrics from a given region after tricorder
// has begun reading variables in that same region to service another
// in-process request, tricorder will skip calling the region’s update
// function for the incoming request. In this case, the two requests will
// read the same data from that region.
type Region region

// NewRegion creates a new region with a particular update function
func RegisterRegion(updateFunc func()) *Region {
	return (*Region)(newRegion(updateFunc))
}

// RegisterMetric registers a single metric with the health system.
// path is the absolute path of the metric e.g "/proc/rpc";
// metric is the metric to register;
// unit is the unit of measurement for the metric;
// description is the description of the metric.
// RegisterMetric returns ErrPathInUse if path already represents a metric
// or a directory.
// RegisterMetric panics if metric is not of a valid type.
func RegisterMetric(
	path string,
	metric interface{},
	unit units.Unit,
	description string) error {
	return root.registerMetric(
		newPathSpec(path), metric, nil, unit, description)
}

// RegisterMetricWithRegion works just like RegisterMetrics but allows
// the caller to specify the region to which the variable or callback function
// being registered belongs. RegisterMetricWithRegion ignores the region
// parameter when registering a distribution.
func RegisterMetricInRegion(
	path string,
	metric interface{},
	r *Region,
	unit units.Unit,
	description string) error {
	return root.registerMetric(newPathSpec(path), metric, (*region)(r), unit, description)
}

// UnregisterPath unregisters the metric or DirectorySpec at the given path.
// UnregisterPath ignores requests to unregister the root path.
func UnregisterPath(path string) {
	root.unregisterPath(newPathSpec(path))
}

// Bucketer represents the organization of values into buckets for
// distributions. Bucketer instances are immutable.
type Bucketer struct {
	pieces []*bucketPiece
}

var (
	// Ranges in powers of two
	PowersOfTwo = NewExponentialBucketer(20, 1.0, 2.0)
	// Ranges in powers of four
	PowersOfFour = NewExponentialBucketer(11, 1.0, 4.0)
	// Ranges in powers of 10
	PowersOfTen = NewExponentialBucketer(7, 1.0, 10.0)
)

// NewExponentialBucketer returns a Bucketer representing buckets on
// a geometric scale. NewExponentialBucketer(25, 3.0, 1.7) means 25 buckets
// starting with <3.0; 3.0 - 5.1; 5.1 - 8.67; 8.67 - 14.739 etc.
// NewExponentialBucketer panics if count < 2 or if start <= 0 or if scale <= 1.
func NewExponentialBucketer(count int, start, scale float64) *Bucketer {
	return newBucketerFromEndpoints(
		newExponentialBucketerStream(count, start, scale))
}

// NewLinearBucketer returns a Bucketer representing bucktes on
// a linear scale. NewLinearBucketer(5, 0, 10) means 5 buckets
// starting with <0; 0-10; 10-20; 20-30; >=30.
// NewLinearBucketer panics if count < 2 or if increment <= 0.
func NewLinearBucketer(count int, start, increment float64) *Bucketer {
	return newBucketerFromEndpoints(
		newLinearBucketerStream(count, start, increment))
}

// NewArbitraryBucketer returns a Bucketer representing specific endpoints
// NewArbitraryBucketer(10.0, 20.0, 30.0) means 4 buckets:
// <10.0; 10.0 - 20.0; 20.0 - 30.0; >= 30.0.
// NewArbitraryBucketer panics if it is called with no arguments.
// It is the caller's responsibility to ensure that the arguments are in
// ascending order.
func NewArbitraryBucketer(endpoints ...float64) *Bucketer {
	return newBucketerFromEndpoints(endpoints)
}

// NewGeometricBucketer returns a Bucketer representing endpoints
// of the form 10^k, 2*10^k, 5*10^k. lower is the lower bound of
// the endpoints; upper is the upper bound of the endpoints.
// NewGeometricBucker(0.5, 50) ==>
// <0.5; 0.5-1; 1-2; 2-5; 5-10; 10-20; 20-50; >=50
func NewGeometricBucketer(lower, upper float64) *Bucketer {
	return newBucketerFromEndpoints(
		newGeometricBucketerStream(lower, upper))
}

// NewCumulativeDistribution creates a new CumulativeDistribution that uses
// this bucketer to distribute values.
func (b *Bucketer) NewCumulativeDistribution() *CumulativeDistribution {
	return (*CumulativeDistribution)(newDistribution(b, false))
}

// NewNonCumulativeDistribution creates a new NonCumulativeDistribution that
// uses this bucketer to distribute values.
func (b *Bucketer) NewNonCumulativeDistribution() *NonCumulativeDistribution {
	return (*NonCumulativeDistribution)(newDistribution(b, true))
}

// CumulativeDistribution represents a metric that is a distribution of
// values. Cumulative distributions only receive new values.
type CumulativeDistribution distribution

// Add adds a single value to this CumulativeDistribution instance.
// value can be a float32, float64, or a time.Duration.
// If a time.Duration, Add converts it to the same unit of time specified in
// the RegisterMetric call made to register this instance.
func (c *CumulativeDistribution) Add(value interface{}) {
	(*distribution)(c).Add(value)
}

// Unlike in CumulativeDistributions,values in NonCumulativeDistributions
// can change shifting from bucket to bucket.
type NonCumulativeDistribution distribution

// Add adds a single value to this NonCumulativeDistribution instance.
// value can be a float32, float64, or a time.Duration.
// If a time.Duration, Add converts it to the same unit of time specified in
// the RegisterMetric call made to register this instance.
func (c *NonCumulativeDistribution) Add(value interface{}) {
	(*distribution)(c).Add(value)
}

// Update updates a value in this NonCumulativeDistribution instance.
// oldValue and newValue can be a float32, float64, or a time.Duration.
// If a time.Duration, Update converts them to the same unit of time specified
// in the RegisterMetric call made to register this instance.
// The reliability of Update() depends on the caller providing the correct
// old value of what is being changed. Failure to do this results in
// undefined behavior.
// Update updates all distribution statistics in the expected way; however,
// it updates min and max such that min only gets smaller and max only gets
// larger. If update Updates values such that they fall into a narrower
// range than before, min and max remain unchanged to indicate the all-time
// min and all-time max. To have min and max reflect the current min and max
// instead of the all-time min and max, see UpdateMinMax().
func (d *NonCumulativeDistribution) Update(oldValue, newValue interface{}) {
	(*distribution)(d).Update(oldValue, newValue)
}

// UpdateMinMax() estimates the current min and max of this distribution.
// and updates min and max accordingly.
// As this call may be expensive, clients need not use unless both are
// true:
// 1) the client has made calls to Update which narrowed the current
// min and max.
// 2) The client wants min and max to reflect the current min and max
// instead of the all-time min and max.
// This method only estimates. The only guarantees that it makes
// upon returning are:
// original_min <= min <= current_min and old_max >= max >= current_max.
// In fact, calling this method may do nothing at all which would still be
// correct behavior.
func (d *NonCumulativeDistribution) UpdateMinMax() {
}

// Remove removes a value from this NonCumulativeDistribution instance.
// valueToBeRemoved can be a float32, float64, or a time.Duration.
// If a time.Duration, Remove converts it to the same unit of time
// specified in the RegisterMetric call made to register this instance.
// The reliability of Remove() depends on the caller providing a value
// already in the distribution. Failure to do this results in
// undefined behavior.
// Remove updates all distribution statistics in the expected way; however,
// it leaves min and max unchanged.
// To have min and max reflect the current min and max
// instead of the all-time min and max, see UpdateMinMax().
func (d *NonCumulativeDistribution) Remove(valueToBeRemoved interface{}) {
	(*distribution)(d).Remove(valueToBeRemoved)
}

// Sum returns the sum of the values in this distribution.
func (d *NonCumulativeDistribution) Sum() float64 {
	return (*distribution)(d).Sum()
}

// DirectorySpec represents a specific directory in the heirarchy of
// metrics.
type DirectorySpec directory

// RegisterDirectory returns the the DirectorySpec registered with path.
// If nothing is registered with path, RegisterDirectory registers a
// new DirectorySpec with path and returns it.
// RegisterDirectory returns ErrPathInUse if path is already associated
// with a metric.
func RegisterDirectory(path string) (dirSpec *DirectorySpec, err error) {
	r, e := root.registerDirectory(newPathSpec(path))
	return (*DirectorySpec)(r), e
}

// RegisterMetric works just like the package level RegisterMetric
// except that path is relative to this DirectorySpec.
func (d *DirectorySpec) RegisterMetric(
	path string,
	metric interface{},
	unit units.Unit,
	description string) error {
	return (*directory)(d).registerMetric(newPathSpec(path), metric, nil, unit, description)
}

// RegisterMetricWithRegion works just like the package level
// RegisterMetricWithRegion except that path is relative to this
// DirectorySpec.
func (d *DirectorySpec) RegisterMetricInRegion(
	path string,
	metric interface{},
	r *Region,
	unit units.Unit,
	description string) error {
	return (*directory)(d).registerMetric(newPathSpec(path), metric, (*region)(r), unit, description)
}

// RegisterDirectory works just like the package level RegisterDirectory
// except that path is relative to this DirectorySpec.
func (d *DirectorySpec) RegisterDirectory(
	path string) (dirSpec *DirectorySpec, err error) {
	r, e := (*directory)(d).registerDirectory(newPathSpec(path))
	return (*DirectorySpec)(r), e
}

// Returns the absolute path this object represents
func (d *DirectorySpec) AbsPath() string {
	return (*directory)(d).AbsPath()
}

// UnregisterPath works just like the package level UnregisterPath
// except that path is relative to this DirectorySpec.
func (d *DirectorySpec) UnregisterPath(path string) {
	(*directory)(d).unregisterPath(newPathSpec(path))
}

// UnregisterDirectory unregisters this DirectorySpec instance along with
// all metrics and directories within it. The caller can unregister any
// DirectorySpec instance except the one representing the top level directory.
// That DirectorySpec instance simply ignores calls to UnregisterDirectory.
// Using an unregistered DirectorySpec instance to register new metrics may
// cause a panic.
func (d *DirectorySpec) UnregisterDirectory() {
	(*directory)(d).unregisterDirectory()
}

// RegisterFlags registers each application flag as a metric under /proc/flags.
func RegisterFlags() {
	registerFlags()
}

// SetFlagUnit sets the unit for a specific flag. If the flag is a
// time.Duration, the default unit is units.Second; otherwise the default
// unit is units.None. If the client wishes to override the default unit for
// a flag, they call this after registering a flag but before calling
// ParseFlags.
func SetFlagUnit(name string, unit units.Unit) {
	setFlagUnit(name, unit)
}
<|endoftext|>


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

package types

import "fmt"

// A Package describes a Go package.
type Package struct {
	path     string
	name     string
	scope    *Scope
	complete bool
	imports  []*Package
	fake     bool // scope lookup errors are silently dropped if package is fake (internal use only)
}

// NewPackage returns a new Package for the given package path,
// name, and scope. The package is not complete and contains no
// explicit imports.
func NewPackage(path, name string, scope *Scope) *Package {
	return &Package{path: path, name: name, scope: scope}
}

// Path returns the package path.
func (pkg *Package) Path() string { return pkg.path }

// Name returns the package name.
func (pkg *Package) Name() string { return pkg.name }

// Scope returns the (complete or incomplete) package scope
// holding the objects declared at package level (TypeNames,
// Consts, Vars, and Funcs).
func (pkg *Package) Scope() *Scope { return pkg.scope }

// A package is complete if its scope contains (at least) all
// exported objects; otherwise it is incomplete.
func (pkg *Package) Complete() bool { return pkg.complete }

// Imports returns the list of packages explicitly imported by
// pkg; the list is in source order. Package unsafe is excluded.
func (pkg *Package) Imports() []*Package { return pkg.imports }

// MarkComplete marks a package as complete.
func (pkg *Package) MarkComplete() { pkg.complete = true }

func (pkg *Package) String() string {
	return fmt.Sprintf("package %s (%s)", pkg.name, pkg.path)
}
<commit_msg>go.tools/go/types: add Package.SetImports<commit_after>// Copyright 2013 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package types

import "fmt"

// A Package describes a Go package.
type Package struct {
	path     string
	name     string
	scope    *Scope
	complete bool
	imports  []*Package
	fake     bool // scope lookup errors are silently dropped if package is fake (internal use only)
}

// NewPackage returns a new Package for the given package path,
// name, and scope. The package is not complete and contains no
// explicit imports.
func NewPackage(path, name string, scope *Scope) *Package {
	return &Package{path: path, name: name, scope: scope}
}

// Path returns the package path.
func (pkg *Package) Path() string { return pkg.path }

// Name returns the package name.
func (pkg *Package) Name() string { return pkg.name }

// Scope returns the (complete or incomplete) package scope
// holding the objects declared at package level (TypeNames,
// Consts, Vars, and Funcs).
func (pkg *Package) Scope() *Scope { return pkg.scope }

// A package is complete if its scope contains (at least) all
// exported objects; otherwise it is incomplete.
func (pkg *Package) Complete() bool { return pkg.complete }

// MarkComplete marks a package as complete.
func (pkg *Package) MarkComplete() { pkg.complete = true }

// Imports returns the list of packages explicitly imported by
// pkg; the list is in source order. Package unsafe is excluded.
func (pkg *Package) Imports() []*Package { return pkg.imports }

// SetImports sets the list of explicitly imported packages to list.
// It is the caller's responsibility to make sure list elements are unique.
func (pkg *Package) SetImports(list []*Package) { pkg.imports = list }

func (pkg *Package) String() string {
	return fmt.Sprintf("package %s (%s)", pkg.name, pkg.path)
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"io"
	"os"
	"os/exec"
	"testing"
)

func TestLoginWithoutTTY(t *testing.T) {
	cmd := exec.Command(dockerBinary, "login")
	// setup STDOUT and STDERR so that we see any output and errors in our console
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr

	// create a buffer with text then a new line as a return
	buf := bytes.NewBuffer([]byte("buffer test string \n"))

	// use a pipe for stdin and manually copy the data so that
	// the process does not get the TTY
	in, err := cmd.StdinPipe()
	if err != nil {
		t.Fatal(err)
	}
	// copy the bytes into the commands stdin along with a new line
	go io.Copy(in, buf)

	// run the command and block until it's done
	if err := cmd.Run(); err == nil {
		t.Fatal("Expected non nil err when loginning in & TTY not available")
	}

	logDone("login - login without TTY")
}
<commit_msg>Suppress output of TestLoginWithoutTTY<commit_after>package main

import (
	"bytes"
	"io"
	"os/exec"
	"testing"
)

func TestLoginWithoutTTY(t *testing.T) {
	cmd := exec.Command(dockerBinary, "login")

	// create a buffer with text then a new line as a return
	buf := bytes.NewBuffer([]byte("buffer test string \n"))

	// use a pipe for stdin and manually copy the data so that
	// the process does not get the TTY
	in, err := cmd.StdinPipe()
	if err != nil {
		t.Fatal(err)
	}
	// copy the bytes into the commands stdin along with a new line
	go io.Copy(in, buf)

	// run the command and block until it's done
	if err := cmd.Run(); err == nil {
		t.Fatal("Expected non nil err when loginning in & TTY not available")
	}

	logDone("login - login without TTY")
}
<|endoftext|>


<commit_before>package handlers

import (
	"net/http"
	"strings"
)

const defaultCookieName = "ghost.sid"

type Session struct {
	ID   string
	Data map[interface{}]interface{}
}

type SessionOptions struct {
	Store          SessionStore
	CookieTemplate *http.Cookie
	TrustProxy     bool
}

type sessResponseWriter struct {
	http.ResponseWriter
	sess *Session
}

func (this *sessResponseWriter) WrappedWriter() http.ResponseWriter {
	return this.ResponseWriter
}

func SessionHandler(h http.Handler, opts *SessionOptions) http.Handler {
	ck := opts.CookieTemplate
	if ck == nil {
		ck = &http.Cookie{
			Name:   defaultCookieName,
			Path:   "/",
			MaxAge: 0,
		}
	}
	return http.HandlerFunc(
		func(w http.ResponseWriter, r *http.Request) {
			if _, ok := GetSession(w); ok {
				// Self-awareness
				h.ServeHTTP(w, r)
				return
			}

			if strings.Index(r.URL.Path, ck.Path) != 0 {
				// Session cookie does not apply to this path
				h.ServeHTTP(w, r)
				return
			}
		})
}

// Helper function to retrieve the session-augmented writer.
func GetSession(w http.ResponseWriter) (*Session, bool) {
	ss, ok := GetResponseWriter(w, func(tst http.ResponseWriter) bool {
		_, ok := tst.(*sessResponseWriter)
		return ok
	})
	if ok {
		return ss.(*sessResponseWriter).sess, true
	}
	return nil, false
}
<commit_msg>more work on session handler<commit_after>package handlers

import (
	"errors"
	"github.com/gorilla/securecookie"
	"net/http"
	"strings"
)

const defaultCookieName = "ghost.sid"

var (
	ErrSessionSecretMissing = errors.New("session secret is missing")
)

type Session struct {
	ID   string
	Data map[interface{}]interface{}
}

type SessionOptions struct {
	Store          SessionStore
	CookieTemplate http.Cookie
	TrustProxy     bool
	Secret         string
}

type sessResponseWriter struct {
	http.ResponseWriter
	sess      *Session
	sessStore SessionStore
	sessSent  bool
}

func (this *sessResponseWriter) WrappedWriter() http.ResponseWriter {
	return this.ResponseWriter
}

func (this *sessResponseWriter) Write(data []byte) (int, error) {
	if !this.sessSent {
		this.sendSessionCookie()
		this.sessSent = true
	}
	return this.ResponseWriter.Write(data)
}

func (this *sessResponseWriter) WriteHeader(code int) {
	if !this.sessSent {
		this.sendSessionCookie()
		this.sessSent = true
	}
	this.ResponseWriter.WriteHeader(code)
}

func (this *sessResponseWriter) sendSessionCookie() {
}

func SessionHandler(h http.Handler, opts SessionOptions) http.Handler {
	// Make sure the required cookie fields are set
	if opts.CookieTemplate.Name == "" {
		opts.CookieTemplate.Name = defaultCookieName
	}
	if opts.CookieTemplate.Path == "" {
		opts.CookieTemplate.Path = "/"
	}
	// Secret is required
	if opts.Secret == "" {
		panic(ErrSessionSecretMissing)
	}
	return http.HandlerFunc(
		func(w http.ResponseWriter, r *http.Request) {
			if _, ok := getSessionWriter(w); ok {
				// Self-awareness
				h.ServeHTTP(w, r)
				return
			}

			if strings.Index(r.URL.Path, opts.CookieTemplate.Path) != 0 {
				// Session cookie does not apply to this path
				h.ServeHTTP(w, r)
				return
			}
			// Get the session cookie
			exCk, err := r.Cookie(opts.CookieTemplate.Name)
			if err != nil {
				// TODO : Generate a new Session
			}
			ckSessId, err := parseSignedCookie(exCk, opts.Secret)
			if err != nil {
				// TODO : Generate a new session
			}
			_ = ckSessId
			srw := &sessResponseWriter{w, nil, opts.Store, false}
			h.ServeHTTP(srw, r)
		})
}

// Helper function to retrieve the session for the current request.
func GetSession(w http.ResponseWriter) (*Session, bool) {
	ss, ok := getSessionWriter(w)
	if ok {
		return ss.sess, true
	}
	return nil, false
}

// Helper function to retrieve the session store
func GetSessionStore(w http.ResponseWriter) (SessionStore, bool) {
	ss, ok := getSessionWriter(w)
	if ok {
		return ss.sessStore, true
	}
	return nil, false
}

// Internal helper function to retrieve the session writer object.
func getSessionWriter(w http.ResponseWriter) (*sessResponseWriter, bool) {
	ss, ok := GetResponseWriter(w, func(tst http.ResponseWriter) bool {
		_, ok := tst.(*sessResponseWriter)
		return ok
	})
	if ok {
		return ss.(*sessResponseWriter), true
	}
	return nil, false
}

func parseSignedCookie(ck *http.Cookie, secret string) (string, error) {
	var val string

	sck := securecookie.New([]byte(secret), nil)
	err := sck.Decode(ck.Name, ck.Value, val)
	if err != nil {
		return "", err
	}
	return val, nil
}
<|endoftext|>


<commit_before>package harvester

import (
	"io"
	"time"

	"golang.org/x/text/encoding"
	"golang.org/x/text/transform"

	"github.com/elastic/libbeat/common/streambuf"
)

// timedReader keeps track of last time bytes have been read from underlying
// reader.
type timedReader struct {
	reader       io.Reader
	lastReadTime time.Time // last time we read some data from input stream
}

// lineReader reads lines from underlying reader, decoding the input stream
// using the configured codec. The reader keeps track of bytes consumed
// from raw input stream for every decoded line.
type lineReader struct {
	rawInput   io.Reader
	codec      encoding.Encoding
	bufferSize int

	nl        []byte
	inBuffer  *streambuf.Buffer
	outBuffer *streambuf.Buffer
	inOffset  int // input buffer read offset
	byteCount int // number of bytes decoded from input buffer into output buffer
	decoder   transform.Transformer
}

const maxConsecutiveEmptyReads = 100

func newTimedReader(reader io.Reader) *timedReader {
	r := &timedReader{
		reader: reader,
	}
	return r
}

func (r *timedReader) Read(p []byte) (int, error) {
	var err error
	n := 0

	for i := maxConsecutiveEmptyReads; i > 0; i-- {
		n, err = r.reader.Read(p)
		if n > 0 {
			r.lastReadTime = time.Now()
			break
		}

		if err != nil {
			break
		}
	}

	return n, err
}

func newLineReader(
	input io.Reader,
	codec encoding.Encoding,
	bufferSize int,
) (*lineReader, error) {
	l := &lineReader{}

	if err := l.init(input, codec, bufferSize); err != nil {
		return nil, err
	}

	return l, nil
}

func (l *lineReader) init(
	input io.Reader,
	codec encoding.Encoding,
	bufferSize int,
) error {
	l.rawInput = input
	l.codec = codec
	l.bufferSize = bufferSize

	l.codec.NewEncoder()
	nl, _, err := transform.Bytes(l.codec.NewEncoder(), []byte{'\n'})
	if err != nil {
		return err
	}

	l.nl = nl
	l.decoder = l.codec.NewDecoder()
	l.inBuffer = streambuf.New(nil)
	l.outBuffer = streambuf.New(nil)
	return nil
}

func (l *lineReader) next() ([]byte, int, error) {
	for {
		// read next 'potential' line from input buffer/reader
		err := l.advance()
		if err != nil {
			return nil, 0, err
		}

		// check last decoded byte really being '\n'
		buf := l.outBuffer.Bytes()
		if buf[len(buf)-1] == '\n' {
			break
		}
	}

	// output buffer contains complete line ending with '\n'. Extract
	// byte slice from buffer and reset output buffer.
	bytes, err := l.outBuffer.Collect(l.outBuffer.Len())
	l.outBuffer.Reset()
	if err != nil {
		panic(err)
	}

	// return and reset consumed bytes count
	sz := l.byteCount
	l.byteCount = 0
	return bytes, sz, nil
}

func (l *lineReader) advance() error {
	var idx int

	// fill inBuffer until '\n' sequence has been found in input buffer
	for {
		idx = l.inBuffer.IndexFrom(l.inOffset, l.nl)
		if idx > 0 {
			break
		}

		// increase search offset to reduce iterations on buffer when looping
		newOffset := l.inBuffer.Len() - len(l.nl)
		if newOffset > l.inOffset {
			l.inOffset = newOffset
		}

		// try to read more bytes into buffer
		n := 0
		buf := make([]byte, l.bufferSize)
		n, err := l.rawInput.Read(buf)
		l.inBuffer.Append(buf[:n])
		if err != nil {
			return err
		}

		// empty read => return buffer error (more bytes required error)
		if n == 0 {
			return streambuf.ErrNoMoreBytes
		}
	}

	// found encoded byte sequence for '\n' in buffer
	// -> decode input sequence into outBuffer
	sz, err := l.decode(idx + len(l.nl))

	// consume transformed bytes from input buffer
	err = l.inBuffer.Advance(sz)
	l.inBuffer.Reset()

	l.inOffset = idx + 1 - sz // continue scanning input buffer from last position + 1
	if l.inOffset < 0 {
		// fix inOffset if '\n' has encoding > 8bits + fill line has been decoded
		l.inOffset = 0
	}

	return err
}

func (l *lineReader) decode(end int) (int, error) {
	var err error
	buffer := make([]byte, 1024)
	inBytes := l.inBuffer.Bytes()
	start := 0

	for start < end {
		var nDst, nSrc int

		nDst, nSrc, err = l.decoder.Transform(buffer, inBytes[start:end], false)
		start += nSrc

		l.outBuffer.Append(buffer[:nDst])

		if err != nil {
			if err == transform.ErrShortDst { // continue transforming
				continue
			}
			break
		}
	}

	l.byteCount += start
	return start, err
}

func (l *lineReader) partial() ([]byte, int, error) {
	sz, err := l.decode(l.inBuffer.Len()) // decode all input buffer
	l.inBuffer.Advance(sz)
	l.inBuffer.Reset()

	l.inOffset -= sz
	if l.inOffset < 0 {
		l.inOffset = 0
	}

	bytes, err := l.outBuffer.Collect(l.outBuffer.Len())
	if err != nil {
		panic(err)
	}
	l.outBuffer.Reset()

	sz = l.byteCount
	l.byteCount = 0
	return bytes, sz, err
}
<commit_msg>fix index check<commit_after>package harvester

import (
	"io"
	"time"

	"golang.org/x/text/encoding"
	"golang.org/x/text/transform"

	"github.com/elastic/libbeat/common/streambuf"
)

// timedReader keeps track of last time bytes have been read from underlying
// reader.
type timedReader struct {
	reader       io.Reader
	lastReadTime time.Time // last time we read some data from input stream
}

// lineReader reads lines from underlying reader, decoding the input stream
// using the configured codec. The reader keeps track of bytes consumed
// from raw input stream for every decoded line.
type lineReader struct {
	rawInput   io.Reader
	codec      encoding.Encoding
	bufferSize int

	nl        []byte
	inBuffer  *streambuf.Buffer
	outBuffer *streambuf.Buffer
	inOffset  int // input buffer read offset
	byteCount int // number of bytes decoded from input buffer into output buffer
	decoder   transform.Transformer
}

const maxConsecutiveEmptyReads = 100

func newTimedReader(reader io.Reader) *timedReader {
	r := &timedReader{
		reader: reader,
	}
	return r
}

func (r *timedReader) Read(p []byte) (int, error) {
	var err error
	n := 0

	for i := maxConsecutiveEmptyReads; i > 0; i-- {
		n, err = r.reader.Read(p)
		if n > 0 {
			r.lastReadTime = time.Now()
			break
		}

		if err != nil {
			break
		}
	}

	return n, err
}

func newLineReader(
	input io.Reader,
	codec encoding.Encoding,
	bufferSize int,
) (*lineReader, error) {
	l := &lineReader{}

	if err := l.init(input, codec, bufferSize); err != nil {
		return nil, err
	}

	return l, nil
}

func (l *lineReader) init(
	input io.Reader,
	codec encoding.Encoding,
	bufferSize int,
) error {
	l.rawInput = input
	l.codec = codec
	l.bufferSize = bufferSize

	l.codec.NewEncoder()
	nl, _, err := transform.Bytes(l.codec.NewEncoder(), []byte{'\n'})
	if err != nil {
		return err
	}

	l.nl = nl
	l.decoder = l.codec.NewDecoder()
	l.inBuffer = streambuf.New(nil)
	l.outBuffer = streambuf.New(nil)
	return nil
}

func (l *lineReader) next() ([]byte, int, error) {
	for {
		// read next 'potential' line from input buffer/reader
		err := l.advance()
		if err != nil {
			return nil, 0, err
		}

		// check last decoded byte really being '\n'
		buf := l.outBuffer.Bytes()
		if buf[len(buf)-1] == '\n' {
			break
		}
	}

	// output buffer contains complete line ending with '\n'. Extract
	// byte slice from buffer and reset output buffer.
	bytes, err := l.outBuffer.Collect(l.outBuffer.Len())
	l.outBuffer.Reset()
	if err != nil {
		panic(err)
	}

	// return and reset consumed bytes count
	sz := l.byteCount
	l.byteCount = 0
	return bytes, sz, nil
}

func (l *lineReader) advance() error {
	var idx int
	var err error

	// fill inBuffer until '\n' sequence has been found in input buffer
	for {
		idx = l.inBuffer.IndexFrom(l.inOffset, l.nl)
		if idx >= 0 {
			break
		}
		if err != nil {
			// if no newline and last read returned error, return error now
			return err
		}

		// increase search offset to reduce iterations on buffer when looping
		newOffset := l.inBuffer.Len() - len(l.nl)
		if newOffset > l.inOffset {
			l.inOffset = newOffset
		}

		// try to read more bytes into buffer
		n := 0
		buf := make([]byte, l.bufferSize)
		n, err := l.rawInput.Read(buf)
		l.inBuffer.Append(buf[:n])

		if n == 0 && err != nil {
			return err
		}

		// empty read => return buffer error (more bytes required error)
		if n == 0 {
			return streambuf.ErrNoMoreBytes
		}
	}

	// found encoded byte sequence for '\n' in buffer
	// -> decode input sequence into outBuffer
	sz, err := l.decode(idx + len(l.nl))

	// consume transformed bytes from input buffer
	err = l.inBuffer.Advance(sz)
	l.inBuffer.Reset()

	l.inOffset = idx + 1 - sz // continue scanning input buffer from last position + 1
	if l.inOffset < 0 {
		// fix inOffset if '\n' has encoding > 8bits + fill line has been decoded
		l.inOffset = 0
	}

	return err
}

func (l *lineReader) decode(end int) (int, error) {
	var err error
	buffer := make([]byte, 1024)
	inBytes := l.inBuffer.Bytes()
	start := 0

	for start < end {
		var nDst, nSrc int

		nDst, nSrc, err = l.decoder.Transform(buffer, inBytes[start:end], false)
		start += nSrc

		l.outBuffer.Append(buffer[:nDst])

		if err != nil {
			if err == transform.ErrShortDst { // continue transforming
				continue
			}
			break
		}
	}

	l.byteCount += start
	return start, err
}

func (l *lineReader) partial() ([]byte, int, error) {
	sz, err := l.decode(l.inBuffer.Len()) // decode all input buffer
	l.inBuffer.Advance(sz)
	l.inBuffer.Reset()

	l.inOffset -= sz
	if l.inOffset < 0 {
		l.inOffset = 0
	}

	bytes, err := l.outBuffer.Collect(l.outBuffer.Len())
	if err != nil {
		panic(err)
	}
	l.outBuffer.Reset()

	sz = l.byteCount
	l.byteCount = 0
	return bytes, sz, err
}
<|endoftext|>


<commit_before>package builtin

import (
	// "log"

	"github.com/wanliu/flow/builtin/confirm"
	. "github.com/wanliu/flow/context"

	config "github.com/wanliu/flow/builtin/config"
)

type Confirm struct {
	TryGetEntities

	expMins float64

	ExpMins <-chan float64

	Confirm <-chan Context
	Cancel  <-chan Context
	Expire  <-chan Context

	Out chan<- ReplyData
}

func NewConfirm() interface{} {
	return new(Confirm)
}

func (c *Confirm) OnExpMins(mins float64) {
	c.expMins = mins
}

func (c Confirm) OnExpire(ctx Context) {
	ctx.SetValue(config.CtxKeyConfirm, nil)
}

func (c *Confirm) OnConfirm(ctx Context) {
	cIn := ctx.Value(config.CtxKeyConfirm)

	if cIn != nil {
		cfm := cIn.(confirm.Data)
		reply := cfm.Confirm(ctx)
		c.Out <- ReplyData{reply, ctx}
	} else {
		reply := ReplyData{"确认操作已经过期", ctx}
		c.Out <- reply
	}
}

func (c *Confirm) OnCancel(ctx Context) {
	cIn := ctx.Value(config.CtxKeyConfirm)

	if cIn != nil {
		cfm := &cIn.(confirm.Data)
		reply := cfm.Cancel(ctx)
		c.Out <- ReplyData{reply, ctx}
	} else {
		reply := ReplyData{"确认操作已经过期", ctx}
		c.Out <- reply
	}
}
<commit_msg>fix confirm<commit_after>package builtin

import (
	// "log"

	"github.com/wanliu/flow/builtin/confirm"
	. "github.com/wanliu/flow/context"

	config "github.com/wanliu/flow/builtin/config"
)

type Confirm struct {
	TryGetEntities

	expMins float64

	ExpMins <-chan float64

	Confirm <-chan Context
	Cancel  <-chan Context
	Expire  <-chan Context

	Out chan<- ReplyData
}

func NewConfirm() interface{} {
	return new(Confirm)
}

func (c *Confirm) OnExpMins(mins float64) {
	c.expMins = mins
}

func (c Confirm) OnExpire(ctx Context) {
	ctx.SetValue(config.CtxKeyConfirm, nil)
}

func (c *Confirm) OnConfirm(ctx Context) {
	cIn := ctx.Value(config.CtxKeyConfirm)

	if cIn != nil {
		cfm := cIn.(confirm.Data)
		reply := cfm.Confirm(ctx)
		c.Out <- ReplyData{reply, ctx}
	} else {
		reply := ReplyData{"确认操作已经过期", ctx}
		c.Out <- reply
	}
}

func (c *Confirm) OnCancel(ctx Context) {
	cIn := ctx.Value(config.CtxKeyConfirm)

	if cIn != nil {
		cfm := cIn.(*confirm.Data)
		reply := cfm.Cancel(ctx)
		c.Out <- ReplyData{reply, ctx}
	} else {
		reply := ReplyData{"确认操作已经过期", ctx}
		c.Out <- reply
	}
}
<|endoftext|>


<commit_before>package hostdb

import (
	"crypto/rand"
	"errors"
	"math/big"
	"strconv"
	"testing"

	"github.com/NebulousLabs/Sia/modules"
	"github.com/NebulousLabs/Sia/types"
)

// fakeAddr returns a modules.NetAddress to be used in a HostEntry. Such
// addresses are needed in order to satisfy the HostDB's "1 host per IP" rule.
func fakeAddr(n uint8) modules.NetAddress {
	return modules.NetAddress("127.0.0." + strconv.Itoa(int(n)) + ":0")
}

// uniformTreeVerification checks that everything makes sense in the tree given
// the number of entries that the tree is supposed to have and also given that
// every entropy has the same weight.
func uniformTreeVerification(hdb *HostDB, numEntries int) error {
	// Check that the weight of the hostTree is what is expected.
	expectedWeight := types.NewCurrency64(uint64(numEntries)).Mul(hdb.hostTree.hostEntry.weight)
	if hdb.hostTree.weight.Cmp(expectedWeight) != 0 {
		return errors.New("expected weight is incorrect")
	}

	// Check that the length of activeHosts and the count of hostTree are
	// consistent.
	if len(hdb.activeHosts) != numEntries {
		return errors.New("unexpected number of active hosts")
	}

	// Select many random hosts and do naive statistical analysis on the
	// results.
	if !testing.Short() {
		// Pull a bunch of random hosts and count how many times we pull each
		// host.
		selectionMap := make(map[modules.NetAddress]int)
		expected := 100
		for i := 0; i < expected*numEntries; i++ {
			entries := hdb.RandomHosts(1)
			if len(entries) == 0 {
				return errors.New("no hosts!")
			}
			selectionMap[entries[0].IPAddress] = selectionMap[entries[0].IPAddress] + 1
		}

		// See if each host was selected enough times.
		errorBound := 64 // Pretty large, but will still detect if something is seriously wrong.
		for _, count := range selectionMap {
			if count < expected-errorBound || count > expected+errorBound {
				return errors.New("error bound was breached")
			}
		}
	}

	// Try removing an re-adding all hosts.
	var removedEntries []*hostEntry
	for {
		if hdb.hostTree.weight.IsZero() {
			break
		}
		randWeight, err := rand.Int(rand.Reader, hdb.hostTree.weight.Big())
		if err != nil {
			break
		}
		node, err := hdb.hostTree.nodeAtWeight(types.NewCurrency(randWeight))
		if err != nil {
			break
		}
		node.removeNode()
		delete(hdb.activeHosts, node.hostEntry.IPAddress)

		// remove the entry from the hostdb so it won't be selected as a
		// repeat.
		removedEntries = append(removedEntries, node.hostEntry)
	}
	for _, entry := range removedEntries {
		hdb.insertNode(entry)
	}
	return nil
}

// TestWeightedList inserts and removes nodes in a semi-random manner and
// verifies that the tree stays consistent through the adjustments.
func TestWeightedList(t *testing.T) {
	if testing.Short() {
		t.SkipNow()
	}

	// Create a hostdb and 3 equal entries to insert.
	hdb := New()

	// Create a bunch of host entries of equal weight.
	firstInsertions := 64
	for i := 0; i < firstInsertions; i++ {
		entry := hostEntry{
			HostSettings: modules.HostSettings{IPAddress: fakeAddr(uint8(i))},
			weight:       types.NewCurrency64(10),
		}
		hdb.insertNode(&entry)
	}
	err := uniformTreeVerification(hdb, firstInsertions)
	if err != nil {
		t.Error(err)
	}

	// Remove a few hosts and check that the tree is still in order.
	removals := 12
	// Keep a map of what we've removed so far.
	removedMap := make(map[uint8]struct{})
	for i := 0; i < removals; i++ {
		// Try numbers until we roll a number that's not been removed yet.
		var randInt uint8
		for {
			randBig, err := rand.Int(rand.Reader, big.NewInt(int64(firstInsertions)))
			if err != nil {
				t.Fatal(err)
			}
			randInt = uint8(randBig.Int64())
			_, exists := removedMap[randInt]
			if !exists {
				break
			}
		}

		// Remove the entry and add it to the list of removed entries
		err := hdb.RemoveHost(fakeAddr(randInt))
		if err != nil {
			t.Fatal(err)
		}
		removedMap[randInt] = struct{}{}
	}
	err = uniformTreeVerification(hdb, firstInsertions-removals)
	if err != nil {
		t.Error(err)
	}

	// Do some more insertions.
	secondInsertions := 64
	for i := firstInsertions; i < firstInsertions+secondInsertions; i++ {
		entry := hostEntry{
			HostSettings: modules.HostSettings{IPAddress: fakeAddr(uint8(i))},
			weight:       types.NewCurrency64(10),
		}
		hdb.insertNode(&entry)
	}
	err = uniformTreeVerification(hdb, firstInsertions-removals+secondInsertions)
	if err != nil {
		t.Error(err)
	}
}

// TestVariedWeights runs broad statistical tests on selecting hosts with
// multiple different weights.
func TestVariedWeights(t *testing.T) {
	if testing.Short() {
		t.SkipNow()
	}
	hdb := New()

	// insert i hosts with the weights 0, 1, ..., i-1. 100e3 selections will be made
	// per weight added to the tree, the total number of selections necessary
	// will be tallied up as hosts are created.
	hostCount := 5
	expectedPerWeight := int(10e3)
	selections := 0
	for i := 0; i < hostCount; i++ {
		entry := hostEntry{
			HostSettings: modules.HostSettings{IPAddress: fakeAddr(uint8(i))},
			weight:       types.NewCurrency64(uint64(i)),
		}
		hdb.insertNode(&entry)
		selections += i * expectedPerWeight
	}

	// Perform many random selections, noting which host was selected each
	// time.
	selectionMap := make(map[string]int)
	for i := 0; i < selections; i++ {
		randEntry := hdb.RandomHosts(1)
		if len(randEntry) == 0 {
			t.Fatal("no hosts!")
		}
		node, exists := hdb.activeHosts[randEntry[0].IPAddress]
		if !exists {
			t.Fatal("can't find randomly selected node in tree")
		}
		selectionMap[node.hostEntry.weight.String()] += 1
	}

	// Check that each host was selected an expected number of times. An error
	// will be reported if the host of 0 weight is ever selected.
	acceptableError := 0.2
	for weight, timesSelected := range selectionMap {
		intWeight, err := strconv.Atoi(weight)
		if err != nil {
			t.Fatal(err)
		}

		expectedSelected := float64(intWeight * expectedPerWeight)
		if float64(expectedSelected)*acceptableError > float64(timesSelected) || float64(expectedSelected)/acceptableError < float64(timesSelected) {
			t.Error("weighted list not selecting in a uniform distribution based on weight")
			t.Error(expectedSelected)
			t.Error(timesSelected)
		}
	}
}

// TestRepeatInsert inserts 2 hosts with the same address.
func TestRepeatInsert(t *testing.T) {
	if testing.Short() {
		t.SkipNow()
	}
	hdb := New()

	entry1 := hostEntry{
		HostSettings: modules.HostSettings{IPAddress: fakeAddr(0)},
		weight:       types.NewCurrency64(1),
	}
	entry2 := entry1
	hdb.insertNode(&entry1)

	entry2.weight = types.NewCurrency64(100)
	hdb.insertNode(&entry2)
	if len(hdb.activeHosts) != 1 {
		t.Error("insterting the same entry twice should result in only 1 entry in the hostdb")
	}
}

// TestRandomHosts probles the RandomHosts function.
func TestRandomHosts(t *testing.T) {
	if testing.Short() {
		t.SkipNow()
	}
	hdb := New()

	// Insert 3 hosts to be selected.
	entry1 := hostEntry{
		HostSettings: modules.HostSettings{IPAddress: fakeAddr(1)},
		weight:       types.NewCurrency64(1),
	}
	entry2 := hostEntry{
		HostSettings: modules.HostSettings{IPAddress: fakeAddr(2)},
		weight:       types.NewCurrency64(2),
	}
	entry3 := hostEntry{
		HostSettings: modules.HostSettings{IPAddress: fakeAddr(3)},
		weight:       types.NewCurrency64(3),
	}
	hdb.insertNode(&entry1)
	hdb.insertNode(&entry2)
	hdb.insertNode(&entry3)

	if len(hdb.activeHosts) != 3 {
		t.Error("wrong number of hosts")
	}
	if hdb.hostTree.weight.Cmp(types.NewCurrency64(6)) != 0 {
		t.Error("unexpected weight at initialization")
		t.Error(hdb.hostTree.weight)
	}

	// Grab 1 random host.
	randHosts := hdb.RandomHosts(1)
	if len(randHosts) != 1 {
		t.Error("didn't get 1 hosts")
	}
	if len(hdb.activeHosts) != 3 {
		t.Error("wrong number of hosts")
	}
	if hdb.hostTree.weight.Cmp(types.NewCurrency64(6)) != 0 {
		t.Error("unexpected weight at initialization")
	}

	// Grab 2 random hosts.
	randHosts = hdb.RandomHosts(2)
	if len(randHosts) != 2 {
		t.Error("didn't get 2 hosts")
	}
	if len(hdb.activeHosts) != 3 {
		t.Error("wrong number of hosts")
	}
	if hdb.hostTree.weight.Cmp(types.NewCurrency64(6)) != 0 {
		t.Error("unexpected weight at initialization")
	}
	if randHosts[0].IPAddress == randHosts[1].IPAddress {
		t.Error("doubled up")
	}

	// Grab 3 random hosts.
	randHosts = hdb.RandomHosts(3)
	if len(randHosts) != 3 {
		t.Error("didn't get 3 hosts")
	}
	if len(hdb.activeHosts) != 3 {
		t.Error("wrong number of hosts")
	}
	if hdb.hostTree.weight.Cmp(types.NewCurrency64(6)) != 0 {
		t.Error("unexpected weight at initialization")
	}
	if randHosts[0].IPAddress == randHosts[1].IPAddress || randHosts[0].IPAddress == randHosts[2].IPAddress || randHosts[1].IPAddress == randHosts[2].IPAddress {
		t.Error("doubled up")
	}

	// Grab 4 random hosts. 3 should be returned.
	randHosts = hdb.RandomHosts(4)
	if len(randHosts) != 3 {
		t.Error("didn't get 3 hosts")
	}
	if len(hdb.activeHosts) != 3 {
		t.Error("wrong number of hosts")
	}
	if hdb.hostTree.weight.Cmp(types.NewCurrency64(6)) != 0 {
		t.Error("unexpected weight at initialization")
	}
	if randHosts[0].IPAddress == randHosts[1].IPAddress || randHosts[0].IPAddress == randHosts[2].IPAddress || randHosts[1].IPAddress == randHosts[2].IPAddress {
		t.Error("doubled up")
	}
}
<commit_msg>don't spawn scanning threads during tests<commit_after>package hostdb

import (
	"crypto/rand"
	"errors"
	"math/big"
	"strconv"
	"testing"

	"github.com/NebulousLabs/Sia/modules"
	"github.com/NebulousLabs/Sia/sync"
	"github.com/NebulousLabs/Sia/types"
)

// fakeAddr returns a modules.NetAddress to be used in a HostEntry. Such
// addresses are needed in order to satisfy the HostDB's "1 host per IP" rule.
func fakeAddr(n uint8) modules.NetAddress {
	return modules.NetAddress("127.0.0." + strconv.Itoa(int(n)) + ":0")
}

// uniformTreeVerification checks that everything makes sense in the tree given
// the number of entries that the tree is supposed to have and also given that
// every entropy has the same weight.
func uniformTreeVerification(hdb *HostDB, numEntries int) error {
	// Check that the weight of the hostTree is what is expected.
	expectedWeight := types.NewCurrency64(uint64(numEntries)).Mul(hdb.hostTree.hostEntry.weight)
	if hdb.hostTree.weight.Cmp(expectedWeight) != 0 {
		return errors.New("expected weight is incorrect")
	}

	// Check that the length of activeHosts and the count of hostTree are
	// consistent.
	if len(hdb.activeHosts) != numEntries {
		return errors.New("unexpected number of active hosts")
	}

	// Select many random hosts and do naive statistical analysis on the
	// results.
	if !testing.Short() {
		// Pull a bunch of random hosts and count how many times we pull each
		// host.
		selectionMap := make(map[modules.NetAddress]int)
		expected := 100
		for i := 0; i < expected*numEntries; i++ {
			entries := hdb.RandomHosts(1)
			if len(entries) == 0 {
				return errors.New("no hosts!")
			}
			selectionMap[entries[0].IPAddress] = selectionMap[entries[0].IPAddress] + 1
		}

		// See if each host was selected enough times.
		errorBound := 64 // Pretty large, but will still detect if something is seriously wrong.
		for _, count := range selectionMap {
			if count < expected-errorBound || count > expected+errorBound {
				return errors.New("error bound was breached")
			}
		}
	}

	// Try removing an re-adding all hosts.
	var removedEntries []*hostEntry
	for {
		if hdb.hostTree.weight.IsZero() {
			break
		}
		randWeight, err := rand.Int(rand.Reader, hdb.hostTree.weight.Big())
		if err != nil {
			break
		}
		node, err := hdb.hostTree.nodeAtWeight(types.NewCurrency(randWeight))
		if err != nil {
			break
		}
		node.removeNode()
		delete(hdb.activeHosts, node.hostEntry.IPAddress)

		// remove the entry from the hostdb so it won't be selected as a
		// repeat.
		removedEntries = append(removedEntries, node.hostEntry)
	}
	for _, entry := range removedEntries {
		hdb.insertNode(entry)
	}
	return nil
}

// TestWeightedList inserts and removes nodes in a semi-random manner and
// verifies that the tree stays consistent through the adjustments.
func TestWeightedList(t *testing.T) {
	if testing.Short() {
		t.SkipNow()
	}

	// Create a hostdb and 3 equal entries to insert.
	hdb := &HostDB{
		activeHosts: make(map[modules.NetAddress]*hostNode),
		allHosts:    make(map[modules.NetAddress]*hostEntry),
		scanPool:    make(chan *hostEntry, scanPoolSize),

		mu: sync.New(modules.SafeMutexDelay, 1),
	}

	// Create a bunch of host entries of equal weight.
	firstInsertions := 64
	for i := 0; i < firstInsertions; i++ {
		entry := hostEntry{
			HostSettings: modules.HostSettings{IPAddress: fakeAddr(uint8(i))},
			weight:       types.NewCurrency64(10),
		}
		hdb.insertNode(&entry)
	}
	err := uniformTreeVerification(hdb, firstInsertions)
	if err != nil {
		t.Error(err)
	}

	// Remove a few hosts and check that the tree is still in order.
	removals := 12
	// Keep a map of what we've removed so far.
	removedMap := make(map[uint8]struct{})
	for i := 0; i < removals; i++ {
		// Try numbers until we roll a number that's not been removed yet.
		var randInt uint8
		for {
			randBig, err := rand.Int(rand.Reader, big.NewInt(int64(firstInsertions)))
			if err != nil {
				t.Fatal(err)
			}
			randInt = uint8(randBig.Int64())
			_, exists := removedMap[randInt]
			if !exists {
				break
			}
		}

		// Remove the entry and add it to the list of removed entries
		err := hdb.RemoveHost(fakeAddr(randInt))
		if err != nil {
			t.Fatal(err)
		}
		removedMap[randInt] = struct{}{}
	}
	err = uniformTreeVerification(hdb, firstInsertions-removals)
	if err != nil {
		t.Error(err)
	}

	// Do some more insertions.
	secondInsertions := 64
	for i := firstInsertions; i < firstInsertions+secondInsertions; i++ {
		entry := hostEntry{
			HostSettings: modules.HostSettings{IPAddress: fakeAddr(uint8(i))},
			weight:       types.NewCurrency64(10),
		}
		hdb.insertNode(&entry)
	}
	err = uniformTreeVerification(hdb, firstInsertions-removals+secondInsertions)
	if err != nil {
		t.Error(err)
	}
}

// TestVariedWeights runs broad statistical tests on selecting hosts with
// multiple different weights.
func TestVariedWeights(t *testing.T) {
	if testing.Short() {
		t.SkipNow()
	}
	hdb := &HostDB{
		activeHosts: make(map[modules.NetAddress]*hostNode),
		allHosts:    make(map[modules.NetAddress]*hostEntry),
		scanPool:    make(chan *hostEntry, scanPoolSize),

		mu: sync.New(modules.SafeMutexDelay, 1),
	}

	// insert i hosts with the weights 0, 1, ..., i-1. 100e3 selections will be made
	// per weight added to the tree, the total number of selections necessary
	// will be tallied up as hosts are created.
	hostCount := 5
	expectedPerWeight := int(10e3)
	selections := 0
	for i := 0; i < hostCount; i++ {
		entry := hostEntry{
			HostSettings: modules.HostSettings{IPAddress: fakeAddr(uint8(i))},
			weight:       types.NewCurrency64(uint64(i)),
		}
		hdb.insertNode(&entry)
		selections += i * expectedPerWeight
	}

	// Perform many random selections, noting which host was selected each
	// time.
	selectionMap := make(map[string]int)
	for i := 0; i < selections; i++ {
		randEntry := hdb.RandomHosts(1)
		if len(randEntry) == 0 {
			t.Fatal("no hosts!")
		}
		node, exists := hdb.activeHosts[randEntry[0].IPAddress]
		if !exists {
			t.Fatal("can't find randomly selected node in tree")
		}
		selectionMap[node.hostEntry.weight.String()] += 1
	}

	// Check that each host was selected an expected number of times. An error
	// will be reported if the host of 0 weight is ever selected.
	acceptableError := 0.2
	for weight, timesSelected := range selectionMap {
		intWeight, err := strconv.Atoi(weight)
		if err != nil {
			t.Fatal(err)
		}

		expectedSelected := float64(intWeight * expectedPerWeight)
		if float64(expectedSelected)*acceptableError > float64(timesSelected) || float64(expectedSelected)/acceptableError < float64(timesSelected) {
			t.Error("weighted list not selecting in a uniform distribution based on weight")
			t.Error(expectedSelected)
			t.Error(timesSelected)
		}
	}
}

// TestRepeatInsert inserts 2 hosts with the same address.
func TestRepeatInsert(t *testing.T) {
	if testing.Short() {
		t.SkipNow()
	}
	hdb := &HostDB{
		activeHosts: make(map[modules.NetAddress]*hostNode),
		allHosts:    make(map[modules.NetAddress]*hostEntry),
		scanPool:    make(chan *hostEntry, scanPoolSize),

		mu: sync.New(modules.SafeMutexDelay, 1),
	}

	entry1 := hostEntry{
		HostSettings: modules.HostSettings{IPAddress: fakeAddr(0)},
		weight:       types.NewCurrency64(1),
	}
	entry2 := entry1
	hdb.insertNode(&entry1)

	entry2.weight = types.NewCurrency64(100)
	hdb.insertNode(&entry2)
	if len(hdb.activeHosts) != 1 {
		t.Error("insterting the same entry twice should result in only 1 entry in the hostdb")
	}
}

// TestRandomHosts probles the RandomHosts function.
func TestRandomHosts(t *testing.T) {
	if testing.Short() {
		t.SkipNow()
	}
	hdb := &HostDB{
		activeHosts: make(map[modules.NetAddress]*hostNode),
		allHosts:    make(map[modules.NetAddress]*hostEntry),
		scanPool:    make(chan *hostEntry, scanPoolSize),

		mu: sync.New(modules.SafeMutexDelay, 1),
	}

	// Insert 3 hosts to be selected.
	entry1 := hostEntry{
		HostSettings: modules.HostSettings{IPAddress: fakeAddr(1)},
		weight:       types.NewCurrency64(1),
	}
	entry2 := hostEntry{
		HostSettings: modules.HostSettings{IPAddress: fakeAddr(2)},
		weight:       types.NewCurrency64(2),
	}
	entry3 := hostEntry{
		HostSettings: modules.HostSettings{IPAddress: fakeAddr(3)},
		weight:       types.NewCurrency64(3),
	}
	hdb.insertNode(&entry1)
	hdb.insertNode(&entry2)
	hdb.insertNode(&entry3)

	if len(hdb.activeHosts) != 3 {
		t.Error("wrong number of hosts")
	}
	if hdb.hostTree.weight.Cmp(types.NewCurrency64(6)) != 0 {
		t.Error("unexpected weight at initialization")
		t.Error(hdb.hostTree.weight)
	}

	// Grab 1 random host.
	randHosts := hdb.RandomHosts(1)
	if len(randHosts) != 1 {
		t.Error("didn't get 1 hosts")
	}
	if len(hdb.activeHosts) != 3 {
		t.Error("wrong number of hosts")
	}
	if hdb.hostTree.weight.Cmp(types.NewCurrency64(6)) != 0 {
		t.Error("unexpected weight at initialization")
	}

	// Grab 2 random hosts.
	randHosts = hdb.RandomHosts(2)
	if len(randHosts) != 2 {
		t.Error("didn't get 2 hosts")
	}
	if len(hdb.activeHosts) != 3 {
		t.Error("wrong number of hosts")
	}
	if hdb.hostTree.weight.Cmp(types.NewCurrency64(6)) != 0 {
		t.Error("unexpected weight at initialization")
	}
	if randHosts[0].IPAddress == randHosts[1].IPAddress {
		t.Error("doubled up")
	}

	// Grab 3 random hosts.
	randHosts = hdb.RandomHosts(3)
	if len(randHosts) != 3 {
		t.Error("didn't get 3 hosts")
	}
	if len(hdb.activeHosts) != 3 {
		t.Error("wrong number of hosts")
	}
	if hdb.hostTree.weight.Cmp(types.NewCurrency64(6)) != 0 {
		t.Error("unexpected weight at initialization")
	}
	if randHosts[0].IPAddress == randHosts[1].IPAddress || randHosts[0].IPAddress == randHosts[2].IPAddress || randHosts[1].IPAddress == randHosts[2].IPAddress {
		t.Error("doubled up")
	}

	// Grab 4 random hosts. 3 should be returned.
	randHosts = hdb.RandomHosts(4)
	if len(randHosts) != 3 {
		t.Error("didn't get 3 hosts")
	}
	if len(hdb.activeHosts) != 3 {
		t.Error("wrong number of hosts")
	}
	if hdb.hostTree.weight.Cmp(types.NewCurrency64(6)) != 0 {
		t.Error("unexpected weight at initialization")
	}
	if randHosts[0].IPAddress == randHosts[1].IPAddress || randHosts[0].IPAddress == randHosts[2].IPAddress || randHosts[1].IPAddress == randHosts[2].IPAddress {
		t.Error("doubled up")
	}
}
<|endoftext|>