<commit_before>package ddb_test

import (
	"fmt"
	"math"
	"strconv"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/dynamodb"
	. "github.com/runtakun/dynamodb-marshaler-go"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("Unmarshal", func() {

	Context("input value", func() {

		var sut sample

		BeforeEach(func() {
			d := map[string]*dynamodb.AttributeValue{
				"str":     &dynamodb.AttributeValue{S: aws.String("foo")},
				"bool":    &dynamodb.AttributeValue{BOOL: aws.Bool(true)},
				"blob":    &dynamodb.AttributeValue{B: []byte{0x1, 0x2, 0x3}},
				"int":     &dynamodb.AttributeValue{N: aws.String(fmt.Sprintf("%d", -42775))},
				"int8":    &dynamodb.AttributeValue{N: aws.String(fmt.Sprintf("%d", -100))},
				"int16":   &dynamodb.AttributeValue{N: aws.String(fmt.Sprintf("%d", -12345))},
				"int32":   &dynamodb.AttributeValue{N: aws.String(fmt.Sprintf("%d", -1073741014))},
				"int64":   &dynamodb.AttributeValue{N: aws.String(fmt.Sprintf("%d", -1112222111222111))},
				"uint":    &dynamodb.AttributeValue{N: aws.String(fmt.Sprintf("%d", 42775))},
				"uint8":   &dynamodb.AttributeValue{N: aws.String(fmt.Sprintf("%d", 111))},
				"uint16":  &dynamodb.AttributeValue{N: aws.String(fmt.Sprintf("%d", 12345))},
				"uint32":  &dynamodb.AttributeValue{N: aws.String(fmt.Sprintf("%d", 1073741014))},
				"uint64":  &dynamodb.AttributeValue{N: aws.String(fmt.Sprintf("%d", 1112222111222111))},
				"float32": &dynamodb.AttributeValue{N: aws.String(strconv.FormatFloat(math.Pi, 'f', -1, 32))},
				"float64": &dynamodb.AttributeValue{N: aws.String(strconv.FormatFloat(math.E, 'f', -1, 64))},
				"slice": &dynamodb.AttributeValue{SS: []*string{
					aws.String("a"),
					aws.String("b"),
					aws.String("c"),
				}},
				"empty_slice": &dynamodb.AttributeValue{NS: []*string{
					aws.String("1"),
					aws.String("2"),
					aws.String("3"),
				}},
			}
			Unmarshal(d, &sut)
		})

		It("should not be nil", func() {
			Expect(sut).NotTo(BeNil())
		})

		It("should be struct which has `Str` column", func() {
			Expect(sut.Str).To(Equal("foo"))
			Expect(sut.Str).NotTo(BeEmpty())
		})

		It("should be struct which has `Bool` column", func() {
			Expect(sut.Bool).To(BeTrue())
			Expect(sut.Bool).NotTo(BeFalse())
		})

		It("should be struct which has `Blob` column", func() {
			Expect(sut.Blob).Should(HaveLen(3))
			Expect(sut.Blob[0]).To(Equal(byte(1)))
			Expect(sut.Blob[1]).To(Equal(byte(2)))
			Expect(sut.Blob[2]).To(Equal(byte(3)))
		})

		It("should be struct which has `Int` column", func() {
			Expect(sut.Int).To(Equal(-42775))
		})

		It("should be struct which has `Int8` column", func() {
			Expect(sut.Int8).To(Equal(int8(-100)))
		})

		It("should be struct which has `Int16` column", func() {
			Expect(sut.Int16).To(Equal(int16(-12345)))
		})

		It("should be struct which has `Int32` column", func() {
			Expect(sut.Int32).To(Equal(int32(-1073741014)))
		})

		It("should be struct which has `Int64` column", func() {
			Expect(sut.Int64).To(Equal(int64(-1112222111222111)))
		})

		It("should be struct which has `Uint` column", func() {
			Expect(sut.Uint).To(Equal(uint(42775)))
		})

		It("should be struct which has `Uint8` column", func() {
			Expect(sut.Uint8).To(Equal(uint8(111)))
		})

		It("should be struct which has `Uint16` column", func() {
			Expect(sut.Uint16).To(Equal(uint16(12345)))
		})

		It("should be struct which has `Uint32` column", func() {
			Expect(sut.Uint32).To(Equal(uint32(1073741014)))
		})

		It("should be struct which has `float32` column", func() {
			Expect(sut.Float32).To(Equal(float32(math.Pi)))
		})

		It("should be struct which has `float64` column", func() {
			Expect(sut.Float64).To(Equal(float64(math.E)))
		})

		It("should be struct which has `slice` column", func() {
			Expect(sut.Slice).Should(HaveLen(3))
			Expect(sut.Slice[0]).To(Equal("a"))
			Expect(sut.Slice[1]).To(Equal("b"))
			Expect(sut.Slice[2]).To(Equal("c"))
		})

		It("should be struct which has `empty_slice` column", func() {
			Expect(sut.EmptySlice).Should(HaveLen(3))
			Expect(sut.EmptySlice[0]).To(Equal(1))
			Expect(sut.EmptySlice[1]).To(Equal(2))
			Expect(sut.EmptySlice[2]).To(Equal(3))
		})

	})

})
<commit_msg>add test for map<commit_after>package ddb_test

import (
	"fmt"
	"math"
	"strconv"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/dynamodb"
	. "github.com/runtakun/dynamodb-marshaler-go"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("Unmarshal", func() {

	Context("input value", func() {

		var sut sample

		BeforeEach(func() {
			d := map[string]*dynamodb.AttributeValue{
				"str":     &dynamodb.AttributeValue{S: aws.String("foo")},
				"bool":    &dynamodb.AttributeValue{BOOL: aws.Bool(true)},
				"blob":    &dynamodb.AttributeValue{B: []byte{0x1, 0x2, 0x3}},
				"int":     &dynamodb.AttributeValue{N: aws.String(fmt.Sprintf("%d", -42775))},
				"int8":    &dynamodb.AttributeValue{N: aws.String(fmt.Sprintf("%d", -100))},
				"int16":   &dynamodb.AttributeValue{N: aws.String(fmt.Sprintf("%d", -12345))},
				"int32":   &dynamodb.AttributeValue{N: aws.String(fmt.Sprintf("%d", -1073741014))},
				"int64":   &dynamodb.AttributeValue{N: aws.String(fmt.Sprintf("%d", -1112222111222111))},
				"uint":    &dynamodb.AttributeValue{N: aws.String(fmt.Sprintf("%d", 42775))},
				"uint8":   &dynamodb.AttributeValue{N: aws.String(fmt.Sprintf("%d", 111))},
				"uint16":  &dynamodb.AttributeValue{N: aws.String(fmt.Sprintf("%d", 12345))},
				"uint32":  &dynamodb.AttributeValue{N: aws.String(fmt.Sprintf("%d", 1073741014))},
				"uint64":  &dynamodb.AttributeValue{N: aws.String(fmt.Sprintf("%d", 1112222111222111))},
				"float32": &dynamodb.AttributeValue{N: aws.String(strconv.FormatFloat(math.Pi, 'f', -1, 32))},
				"float64": &dynamodb.AttributeValue{N: aws.String(strconv.FormatFloat(math.E, 'f', -1, 64))},
				"slice": &dynamodb.AttributeValue{SS: []*string{
					aws.String("a"),
					aws.String("b"),
					aws.String("c"),
				}},
				"empty_slice": &dynamodb.AttributeValue{NS: []*string{
					aws.String("1"),
					aws.String("2"),
					aws.String("3"),
				}},
				"map": &dynamodb.AttributeValue{M: map[string]*dynamodb.AttributeValue{
					"map_foo": &dynamodb.AttributeValue{S: aws.String("map_bar")},
					"map_int": &dynamodb.AttributeValue{N: aws.String("54321")},
				}},
			}
			Unmarshal(d, &sut)
		})

		It("should not be nil", func() {
			Expect(sut).NotTo(BeNil())
		})

		It("should be struct which has `Str` column", func() {
			Expect(sut.Str).To(Equal("foo"))
			Expect(sut.Str).NotTo(BeEmpty())
		})

		It("should be struct which has `Bool` column", func() {
			Expect(sut.Bool).To(BeTrue())
			Expect(sut.Bool).NotTo(BeFalse())
		})

		It("should be struct which has `Blob` column", func() {
			Expect(sut.Blob).Should(HaveLen(3))
			Expect(sut.Blob[0]).To(Equal(byte(1)))
			Expect(sut.Blob[1]).To(Equal(byte(2)))
			Expect(sut.Blob[2]).To(Equal(byte(3)))
		})

		It("should be struct which has `Int` column", func() {
			Expect(sut.Int).To(Equal(-42775))
		})

		It("should be struct which has `Int8` column", func() {
			Expect(sut.Int8).To(Equal(int8(-100)))
		})

		It("should be struct which has `Int16` column", func() {
			Expect(sut.Int16).To(Equal(int16(-12345)))
		})

		It("should be struct which has `Int32` column", func() {
			Expect(sut.Int32).To(Equal(int32(-1073741014)))
		})

		It("should be struct which has `Int64` column", func() {
			Expect(sut.Int64).To(Equal(int64(-1112222111222111)))
		})

		It("should be struct which has `Uint` column", func() {
			Expect(sut.Uint).To(Equal(uint(42775)))
		})

		It("should be struct which has `Uint8` column", func() {
			Expect(sut.Uint8).To(Equal(uint8(111)))
		})

		It("should be struct which has `Uint16` column", func() {
			Expect(sut.Uint16).To(Equal(uint16(12345)))
		})

		It("should be struct which has `Uint32` column", func() {
			Expect(sut.Uint32).To(Equal(uint32(1073741014)))
		})

		It("should be struct which has `float32` column", func() {
			Expect(sut.Float32).To(Equal(float32(math.Pi)))
		})

		It("should be struct which has `float64` column", func() {
			Expect(sut.Float64).To(Equal(float64(math.E)))
		})

		It("should be struct which has `slice` column", func() {
			Expect(sut.Slice).Should(HaveLen(3))
			Expect(sut.Slice[0]).To(Equal("a"))
			Expect(sut.Slice[1]).To(Equal("b"))
			Expect(sut.Slice[2]).To(Equal("c"))
		})

		It("should be struct which has `empty_slice` column", func() {
			Expect(sut.EmptySlice).Should(HaveLen(3))
			Expect(sut.EmptySlice[0]).To(Equal(1))
			Expect(sut.EmptySlice[1]).To(Equal(2))
			Expect(sut.EmptySlice[2]).To(Equal(3))
		})

		It("should be struct which has `map` column", func() {
			Expect(sut.Map).ShouldNot(BeNil())

			v1, ok1 := sut.Map["map_foo"]
			Expect(ok1).Should(BeTrue())
			Expect(v1).To(Equal("map_bar"))

			v2, ok2 := sut.Map["map_int"]
			Expect(ok2).Should(BeTrue())
			Expect(v2).To(Equal(54321))
		})
	})
})
<|endoftext|>


<commit_before>package s3

import (
	"os"
	"path/filepath"

	"github.com/goamz/goamz/aws"
	"github.com/goamz/goamz/s3"
)

// Uploader for S3
type Uploader struct {
	BaseDir string
}

var bucket *s3.Bucket

// Upload copies a file int a bucket in S3
func (u *Uploader) Upload(source string, destination string, contType string) error {
	reader, err := os.Open(source)
	if err != nil {
		return err
	}
	defer reader.Close()

	var stat os.FileInfo
	stat, err = os.Stat(source)
	if err != nil {
		return err
	}
	size := stat.Size()

	err = bucket.PutReader(destination, reader, size, contType, s3.PublicRead, s3.Options{})
	return err
}

func (u *Uploader) ListDirectory(directory string) ([]string, error) {
	var names []string
	prefix := directory
	delim := ""
	marker := ""
	max := 0
	resp, err := bucket.List(prefix, delim, marker, max)
	if err == nil {
		entries := resp.Contents
		for _, entry := range entries {
			name := filepath.Base(entry.Key)
			names = append(names, name)
		}
	}
	return names, err
}

// CreateDirectory does nothing since a directory does not need to be created on S3
// Directories are virtual, and defined by the path of the object
func (u *Uploader) CreateDirectory(path string) error {
	return nil
}

func Initialize(accessKey string, secretKey string, bucketName string) {
	bucket = retrieveBucket(accessKey, secretKey, bucketName)
}

func retrieveBucket(accessKey, secretKey, bucketName string) *s3.Bucket {
	auth := aws.Auth{AccessKey: accessKey, SecretKey: secretKey}
	client := s3.New(auth, aws.USEast)
	return client.Bucket(bucketName)
}
<commit_msg>s3 uploader has explicit timeouts<commit_after>package s3

import (
	"os"
	"path/filepath"
	"time"

	"github.com/goamz/goamz/aws"
	"github.com/goamz/goamz/s3"
)

// Uploader for S3
type Uploader struct {
	BaseDir string
}

var bucket *s3.Bucket

// Upload copies a file int a bucket in S3
func (u *Uploader) Upload(source string, destination string, contType string) error {
	reader, err := os.Open(source)
	if err != nil {
		return err
	}
	defer reader.Close()

	var stat os.FileInfo
	stat, err = os.Stat(source)
	if err != nil {
		return err
	}
	size := stat.Size()

	err = bucket.PutReader(destination, reader, size, contType, s3.PublicRead, s3.Options{})
	return err
}

func (u *Uploader) ListDirectory(directory string) ([]string, error) {
	var names []string
	prefix := directory
	delim := ""
	marker := ""
	max := 0
	resp, err := bucket.List(prefix, delim, marker, max)
	if err == nil {
		entries := resp.Contents
		for _, entry := range entries {
			name := filepath.Base(entry.Key)
			names = append(names, name)
		}
	}
	return names, err
}

// CreateDirectory does nothing since a directory does not need to be created on S3
// Directories are virtual, and defined by the path of the object
func (u *Uploader) CreateDirectory(path string) error {
	return nil
}

func Initialize(accessKey string, secretKey string, bucketName string) {
	bucket = retrieveBucket(accessKey, secretKey, bucketName)
}

func retrieveBucket(accessKey, secretKey, bucketName string) *s3.Bucket {
	client := s3.S3{
		Auth:   aws.Auth{AccessKey: accessKey, SecretKey: secretKey},
		Region: aws.USEast,
		AttemptStrategy: aws.AttemptStrategy{
			Min:   2,
			Total: 5 * time.Second,
			Delay: 200 * time.Millisecond,
		},
		ConnectTimeout: 1 * time.Second,
		ReadTimeout:    5 * time.Second,
		RequestTimeout: 5 * time.Second,
	}
	return client.Bucket(bucketName)
}
<|endoftext|>


<commit_before>package service

import (
	"context"
	"strings"
	"time"

	"github.com/micro/go-micro/v2/auth"
	"github.com/micro/go-micro/v2/auth/rules"
	pb "github.com/micro/go-micro/v2/auth/service/proto"
	"github.com/micro/go-micro/v2/auth/token"
	"github.com/micro/go-micro/v2/auth/token/jwt"
	"github.com/micro/go-micro/v2/client"
)

// svc is the service implementation of the Auth interface
type svc struct {
	options auth.Options
	auth    pb.AuthService
	rules   pb.RulesService
	jwt     token.Provider
}

func (s *svc) String() string {
	return "service"
}

func (s *svc) Init(opts ...auth.Option) {
	for _, o := range opts {
		o(&s.options)
	}

	if s.options.Client == nil {
		s.options.Client = client.DefaultClient
	}

	s.auth = pb.NewAuthService("go.micro.auth", s.options.Client)
	s.rules = pb.NewRulesService("go.micro.auth", s.options.Client)

	// if we have a JWT public key passed as an option,
	// we can decode tokens with the type "JWT" locally
	// and not have to make an RPC call
	if key := s.options.PublicKey; len(key) > 0 {
		s.jwt = jwt.NewTokenProvider(token.WithPublicKey(key))
	}
}

func (s *svc) Options() auth.Options {
	return s.options
}

// Generate a new account
func (s *svc) Generate(id string, opts ...auth.GenerateOption) (*auth.Account, error) {
	options := auth.NewGenerateOptions(opts...)

	rsp, err := s.auth.Generate(context.TODO(), &pb.GenerateRequest{
		Id:       id,
		Type:     options.Type,
		Secret:   options.Secret,
		Scopes:   options.Scopes,
		Metadata: options.Metadata,
		Provider: options.Provider,
	})
	if err != nil {
		return nil, err
	}

	return serializeAccount(rsp.Account), nil
}

// Grant access to a resource
func (s *svc) Grant(rule *auth.Rule) error {
	access := pb.Access_UNKNOWN
	if rule.Access == auth.AccessGranted {
		access = pb.Access_GRANTED
	} else if rule.Access == auth.AccessDenied {
		access = pb.Access_DENIED
	}

	_, err := s.rules.Create(context.TODO(), &pb.CreateRequest{
		Rule: &pb.Rule{
			Id:       rule.ID,
			Scope:    rule.Scope,
			Priority: rule.Priority,
			Access:   access,
			Resource: &pb.Resource{
				Type:     rule.Resource.Type,
				Name:     rule.Resource.Name,
				Endpoint: rule.Resource.Endpoint,
			},
		},
	})

	return err
}

// Revoke access to a resource
func (s *svc) Revoke(rule *auth.Rule) error {
	_, err := s.rules.Delete(context.TODO(), &pb.DeleteRequest{
		Id: rule.ID,
	})

	return err
}

func (s *svc) Rules(opts ...auth.RulesOption) ([]*auth.Rule, error) {
	var options auth.RulesOptions
	for _, o := range opts {
		o(&options)
	}
	if options.Context == nil {
		options.Context = context.TODO()
	}

	rsp, err := s.rules.List(options.Context, &pb.ListRequest{})
	if err != nil {
		return nil, err
	}

	rules := make([]*auth.Rule, len(rsp.Rules))
	for i, r := range rsp.Rules {
		rules[i] = serializeRule(r)
	}

	return rules, nil
}

// Verify an account has access to a resource
func (s *svc) Verify(acc *auth.Account, res *auth.Resource, opts ...auth.VerifyOption) error {
	var options auth.VerifyOptions
	for _, o := range opts {
		o(&options)
	}

	rs, err := s.Rules(auth.RulesContext(options.Context))
	if err != nil {
		return err
	}

	return rules.Verify(rs, acc, res)
}

// Inspect a token
func (s *svc) Inspect(token string) (*auth.Account, error) {
	// try to decode JWT locally and fall back to srv if an error occurs
	if len(strings.Split(token, ".")) == 3 && s.jwt != nil {
		return s.jwt.Inspect(token)
	}

	// the token is not a JWT or we do not have the keys to decode it,
	// fall back to the auth service
	rsp, err := s.auth.Inspect(context.TODO(), &pb.InspectRequest{Token: token})
	if err != nil {
		return nil, err
	}
	return serializeAccount(rsp.Account), nil
}

// Token generation using an account ID and secret
func (s *svc) Token(opts ...auth.TokenOption) (*auth.Token, error) {
	options := auth.NewTokenOptions(opts...)

	rsp, err := s.auth.Token(context.Background(), &pb.TokenRequest{
		Id:           options.ID,
		Secret:       options.Secret,
		RefreshToken: options.RefreshToken,
		TokenExpiry:  int64(options.Expiry.Seconds()),
	})
	if err != nil {
		return nil, err
	}

	return serializeToken(rsp.Token), nil
}

func serializeToken(t *pb.Token) *auth.Token {
	return &auth.Token{
		AccessToken:  t.AccessToken,
		RefreshToken: t.RefreshToken,
		Created:      time.Unix(t.Created, 0),
		Expiry:       time.Unix(t.Expiry, 0),
	}
}

func serializeAccount(a *pb.Account) *auth.Account {
	return &auth.Account{
		ID:       a.Id,
		Secret:   a.Secret,
		Issuer:   a.Issuer,
		Metadata: a.Metadata,
		Scopes:   a.Scopes,
	}
}

func serializeRule(r *pb.Rule) *auth.Rule {
	var access auth.Access
	if r.Access == pb.Access_GRANTED {
		access = auth.AccessGranted
	} else {
		access = auth.AccessDenied
	}

	return &auth.Rule{
		ID:       r.Id,
		Scope:    r.Scope,
		Access:   access,
		Priority: r.Priority,
		Resource: &auth.Resource{
			Type:     r.Resource.Type,
			Name:     r.Resource.Name,
			Endpoint: r.Resource.Endpoint,
		},
	}
}

// NewAuth returns a new instance of the Auth service
func NewAuth(opts ...auth.Option) auth.Auth {
	options := auth.NewOptions(opts...)
	if options.Client == nil {
		options.Client = client.DefaultClient
	}

	return &svc{
		auth:    pb.NewAuthService("go.micro.auth", options.Client),
		rules:   pb.NewRulesService("go.micro.auth", options.Client),
		options: options,
	}
}
<commit_msg>Cache Rules<commit_after>package service

import (
	"context"
	"strings"
	"time"

	"github.com/micro/go-micro/v2/auth"
	"github.com/micro/go-micro/v2/auth/rules"
	pb "github.com/micro/go-micro/v2/auth/service/proto"
	"github.com/micro/go-micro/v2/auth/token"
	"github.com/micro/go-micro/v2/auth/token/jwt"
	"github.com/micro/go-micro/v2/client"
)

// svc is the service implementation of the Auth interface
type svc struct {
	options auth.Options
	auth    pb.AuthService
	rules   pb.RulesService
	jwt     token.Provider
}

func (s *svc) String() string {
	return "service"
}

func (s *svc) Init(opts ...auth.Option) {
	for _, o := range opts {
		o(&s.options)
	}

	if s.options.Client == nil {
		s.options.Client = client.DefaultClient
	}

	s.auth = pb.NewAuthService("go.micro.auth", s.options.Client)
	s.rules = pb.NewRulesService("go.micro.auth", s.options.Client)

	// if we have a JWT public key passed as an option,
	// we can decode tokens with the type "JWT" locally
	// and not have to make an RPC call
	if key := s.options.PublicKey; len(key) > 0 {
		s.jwt = jwt.NewTokenProvider(token.WithPublicKey(key))
	}
}

func (s *svc) Options() auth.Options {
	return s.options
}

// Generate a new account
func (s *svc) Generate(id string, opts ...auth.GenerateOption) (*auth.Account, error) {
	options := auth.NewGenerateOptions(opts...)

	rsp, err := s.auth.Generate(context.TODO(), &pb.GenerateRequest{
		Id:       id,
		Type:     options.Type,
		Secret:   options.Secret,
		Scopes:   options.Scopes,
		Metadata: options.Metadata,
		Provider: options.Provider,
	})
	if err != nil {
		return nil, err
	}

	return serializeAccount(rsp.Account), nil
}

// Grant access to a resource
func (s *svc) Grant(rule *auth.Rule) error {
	access := pb.Access_UNKNOWN
	if rule.Access == auth.AccessGranted {
		access = pb.Access_GRANTED
	} else if rule.Access == auth.AccessDenied {
		access = pb.Access_DENIED
	}

	_, err := s.rules.Create(context.TODO(), &pb.CreateRequest{
		Rule: &pb.Rule{
			Id:       rule.ID,
			Scope:    rule.Scope,
			Priority: rule.Priority,
			Access:   access,
			Resource: &pb.Resource{
				Type:     rule.Resource.Type,
				Name:     rule.Resource.Name,
				Endpoint: rule.Resource.Endpoint,
			},
		},
	})

	return err
}

// Revoke access to a resource
func (s *svc) Revoke(rule *auth.Rule) error {
	_, err := s.rules.Delete(context.TODO(), &pb.DeleteRequest{
		Id: rule.ID,
	})

	return err
}

func (s *svc) Rules(opts ...auth.RulesOption) ([]*auth.Rule, error) {
	var options auth.RulesOptions
	for _, o := range opts {
		o(&options)
	}
	if options.Context == nil {
		options.Context = context.TODO()
	}

	rsp, err := s.rules.List(options.Context, &pb.ListRequest{}, client.WithCache(time.Second*30))
	if err != nil {
		return nil, err
	}

	rules := make([]*auth.Rule, len(rsp.Rules))
	for i, r := range rsp.Rules {
		rules[i] = serializeRule(r)
	}

	return rules, nil
}

// Verify an account has access to a resource
func (s *svc) Verify(acc *auth.Account, res *auth.Resource, opts ...auth.VerifyOption) error {
	var options auth.VerifyOptions
	for _, o := range opts {
		o(&options)
	}

	rs, err := s.Rules(auth.RulesContext(options.Context))
	if err != nil {
		return err
	}

	return rules.Verify(rs, acc, res)
}

// Inspect a token
func (s *svc) Inspect(token string) (*auth.Account, error) {
	// try to decode JWT locally and fall back to srv if an error occurs
	if len(strings.Split(token, ".")) == 3 && s.jwt != nil {
		return s.jwt.Inspect(token)
	}

	// the token is not a JWT or we do not have the keys to decode it,
	// fall back to the auth service
	rsp, err := s.auth.Inspect(context.TODO(), &pb.InspectRequest{Token: token})
	if err != nil {
		return nil, err
	}
	return serializeAccount(rsp.Account), nil
}

// Token generation using an account ID and secret
func (s *svc) Token(opts ...auth.TokenOption) (*auth.Token, error) {
	options := auth.NewTokenOptions(opts...)

	rsp, err := s.auth.Token(context.Background(), &pb.TokenRequest{
		Id:           options.ID,
		Secret:       options.Secret,
		RefreshToken: options.RefreshToken,
		TokenExpiry:  int64(options.Expiry.Seconds()),
	})
	if err != nil {
		return nil, err
	}

	return serializeToken(rsp.Token), nil
}

func serializeToken(t *pb.Token) *auth.Token {
	return &auth.Token{
		AccessToken:  t.AccessToken,
		RefreshToken: t.RefreshToken,
		Created:      time.Unix(t.Created, 0),
		Expiry:       time.Unix(t.Expiry, 0),
	}
}

func serializeAccount(a *pb.Account) *auth.Account {
	return &auth.Account{
		ID:       a.Id,
		Secret:   a.Secret,
		Issuer:   a.Issuer,
		Metadata: a.Metadata,
		Scopes:   a.Scopes,
	}
}

func serializeRule(r *pb.Rule) *auth.Rule {
	var access auth.Access
	if r.Access == pb.Access_GRANTED {
		access = auth.AccessGranted
	} else {
		access = auth.AccessDenied
	}

	return &auth.Rule{
		ID:       r.Id,
		Scope:    r.Scope,
		Access:   access,
		Priority: r.Priority,
		Resource: &auth.Resource{
			Type:     r.Resource.Type,
			Name:     r.Resource.Name,
			Endpoint: r.Resource.Endpoint,
		},
	}
}

// NewAuth returns a new instance of the Auth service
func NewAuth(opts ...auth.Option) auth.Auth {
	options := auth.NewOptions(opts...)
	if options.Client == nil {
		options.Client = client.DefaultClient
	}

	return &svc{
		auth:    pb.NewAuthService("go.micro.auth", options.Client),
		rules:   pb.NewRulesService("go.micro.auth", options.Client),
		options: options,
	}
}
<|endoftext|>


<commit_before>package http

import (
	"context"
	"fmt"
	"io/ioutil"
	"net/url"
	"os"
	"path/filepath"

	"github.com/go-git/go-git/v5/plumbing"
	"github.com/go-git/go-git/v5/plumbing/protocol/packp"
	"github.com/go-git/go-git/v5/plumbing/transport"
	"github.com/go-git/go-git/v5/plumbing/transport/test"

	fixtures "github.com/go-git/go-git-fixtures/v4"
	. "gopkg.in/check.v1"
)

type UploadPackSuite struct {
	test.UploadPackSuite
	BaseSuite
}

var _ = Suite(&UploadPackSuite{})

func (s *UploadPackSuite) SetUpSuite(c *C) {
	s.BaseSuite.SetUpTest(c)
	s.UploadPackSuite.Client = DefaultClient
	s.UploadPackSuite.Endpoint = s.prepareRepository(c, fixtures.Basic().One(), "basic.git")
	s.UploadPackSuite.EmptyEndpoint = s.prepareRepository(c, fixtures.ByTag("empty").One(), "empty.git")
	s.UploadPackSuite.NonExistentEndpoint = s.newEndpoint(c, "non-existent.git")
}

// Overwritten, different behaviour for HTTP.
func (s *UploadPackSuite) TestAdvertisedReferencesNotExists(c *C) {
	r, err := s.Client.NewUploadPackSession(s.NonExistentEndpoint, s.EmptyAuth)
	c.Assert(err, IsNil)
	info, err := r.AdvertisedReferences()
	c.Assert(err, Equals, transport.ErrRepositoryNotFound)
	c.Assert(info, IsNil)
}

func (s *UploadPackSuite) TestuploadPackRequestToReader(c *C) {
	r := packp.NewUploadPackRequest()
	r.Wants = append(r.Wants, plumbing.NewHash("d82f291cde9987322c8a0c81a325e1ba6159684c"))
	r.Wants = append(r.Wants, plumbing.NewHash("2b41ef280fdb67a9b250678686a0c3e03b0a9989"))
	r.Haves = append(r.Haves, plumbing.NewHash("6ecf0ef2c2dffb796033e5a02219af86ec6584e5"))

	sr, err := uploadPackRequestToReader(r)
	c.Assert(err, IsNil)
	b, _ := ioutil.ReadAll(sr)
	c.Assert(string(b), Equals,
		"0032want 2b41ef280fdb67a9b250678686a0c3e03b0a9989\n"+
			"0032want d82f291cde9987322c8a0c81a325e1ba6159684c\n0000"+
			"0032have 6ecf0ef2c2dffb796033e5a02219af86ec6584e5\n"+
			"0009done\n",
	)
}

func (s *UploadPackSuite) prepareRepository(c *C, f *fixtures.Fixture, name string) *transport.Endpoint {
	fs := f.DotGit()

	err := fixtures.EnsureIsBare(fs)
	c.Assert(err, IsNil)

	path := filepath.Join(s.base, name)
	err = os.Rename(fs.Root(), path)
	c.Assert(err, IsNil)

	return s.newEndpoint(c, name)
}

func (s *UploadPackSuite) newEndpoint(c *C, name string) *transport.Endpoint {
	ep, err := transport.NewEndpoint(fmt.Sprintf("http://localhost:%d/%s", s.port, name))
	c.Assert(err, IsNil)

	return ep
}

func (s *UploadPackSuite) TestAdvertisedReferencesRedirectPath(c *C) {
	endpoint, _ := transport.NewEndpoint("https://gitlab.com/gitlab-org/gitter/webapp")

	session, err := s.Client.NewUploadPackSession(endpoint, s.EmptyAuth)
	c.Assert(err, IsNil)

	info, err := session.AdvertisedReferences()
	c.Assert(err, IsNil)
	c.Assert(info, NotNil)

	url := session.(*upSession).endpoint.String()
	c.Assert(url, Equals, "https://gitlab.com/gitlab-org/gitter/webapp.git")
}

func (s *UploadPackSuite) TestAdvertisedReferencesRedirectSchema(c *C) {
	endpoint, _ := transport.NewEndpoint("http://github.com/git-fixtures/basic")

	session, err := s.Client.NewUploadPackSession(endpoint, s.EmptyAuth)
	c.Assert(err, IsNil)

	info, err := session.AdvertisedReferences()
	c.Assert(err, IsNil)
	c.Assert(info, NotNil)

	url := session.(*upSession).endpoint.String()
	c.Assert(url, Equals, "https://github.com/git-fixtures/basic")
}

func (s *UploadPackSuite) TestAdvertisedReferencesContext(c *C) {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	endpoint, _ := transport.NewEndpoint("http://github.com/git-fixtures/basic")

	session, err := s.Client.NewUploadPackSession(endpoint, s.EmptyAuth)
	c.Assert(err, IsNil)

	info, err := session.AdvertisedReferencesContext(ctx)
	c.Assert(err, IsNil)
	c.Assert(info, NotNil)

	url := session.(*upSession).endpoint.String()
	c.Assert(url, Equals, "https://github.com/git-fixtures/basic")
}

func (s *UploadPackSuite) TestAdvertisedReferencesContextCanceled(c *C) {
	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	endpoint, _ := transport.NewEndpoint("http://github.com/git-fixtures/basic")

	session, err := s.Client.NewUploadPackSession(endpoint, s.EmptyAuth)
	c.Assert(err, IsNil)

	info, err := session.AdvertisedReferencesContext(ctx)
	c.Assert(err, DeepEquals, &url.Error{Op: "Get", URL: "http://github.com/git-fixtures/basic/info/refs?service=git-upload-pack", Err: context.Canceled})
	c.Assert(info, IsNil)
}
<commit_msg>plumbing: transport/http, skip flaky test<commit_after>package http

import (
	"context"
	"fmt"
	"io/ioutil"
	"net/url"
	"os"
	"path/filepath"

	"github.com/go-git/go-git/v5/plumbing"
	"github.com/go-git/go-git/v5/plumbing/protocol/packp"
	"github.com/go-git/go-git/v5/plumbing/transport"
	"github.com/go-git/go-git/v5/plumbing/transport/test"

	fixtures "github.com/go-git/go-git-fixtures/v4"
	. "gopkg.in/check.v1"
)

type UploadPackSuite struct {
	test.UploadPackSuite
	BaseSuite
}

var _ = Suite(&UploadPackSuite{})

func (s *UploadPackSuite) SetUpSuite(c *C) {
	s.BaseSuite.SetUpTest(c)
	s.UploadPackSuite.Client = DefaultClient
	s.UploadPackSuite.Endpoint = s.prepareRepository(c, fixtures.Basic().One(), "basic.git")
	s.UploadPackSuite.EmptyEndpoint = s.prepareRepository(c, fixtures.ByTag("empty").One(), "empty.git")
	s.UploadPackSuite.NonExistentEndpoint = s.newEndpoint(c, "non-existent.git")
}

// Overwritten, different behaviour for HTTP.
func (s *UploadPackSuite) TestAdvertisedReferencesNotExists(c *C) {
	r, err := s.Client.NewUploadPackSession(s.NonExistentEndpoint, s.EmptyAuth)
	c.Assert(err, IsNil)
	info, err := r.AdvertisedReferences()
	c.Assert(err, Equals, transport.ErrRepositoryNotFound)
	c.Assert(info, IsNil)
}

func (s *UploadPackSuite) TestuploadPackRequestToReader(c *C) {
	r := packp.NewUploadPackRequest()
	r.Wants = append(r.Wants, plumbing.NewHash("d82f291cde9987322c8a0c81a325e1ba6159684c"))
	r.Wants = append(r.Wants, plumbing.NewHash("2b41ef280fdb67a9b250678686a0c3e03b0a9989"))
	r.Haves = append(r.Haves, plumbing.NewHash("6ecf0ef2c2dffb796033e5a02219af86ec6584e5"))

	sr, err := uploadPackRequestToReader(r)
	c.Assert(err, IsNil)
	b, _ := ioutil.ReadAll(sr)
	c.Assert(string(b), Equals,
		"0032want 2b41ef280fdb67a9b250678686a0c3e03b0a9989\n"+
			"0032want d82f291cde9987322c8a0c81a325e1ba6159684c\n0000"+
			"0032have 6ecf0ef2c2dffb796033e5a02219af86ec6584e5\n"+
			"0009done\n",
	)
}

func (s *UploadPackSuite) prepareRepository(c *C, f *fixtures.Fixture, name string) *transport.Endpoint {
	fs := f.DotGit()

	err := fixtures.EnsureIsBare(fs)
	c.Assert(err, IsNil)

	path := filepath.Join(s.base, name)
	err = os.Rename(fs.Root(), path)
	c.Assert(err, IsNil)

	return s.newEndpoint(c, name)
}

func (s *UploadPackSuite) newEndpoint(c *C, name string) *transport.Endpoint {
	ep, err := transport.NewEndpoint(fmt.Sprintf("http://localhost:%d/%s", s.port, name))
	c.Assert(err, IsNil)

	return ep
}

func (s *UploadPackSuite) TestAdvertisedReferencesRedirectPath(c *C) {
	endpoint, _ := transport.NewEndpoint("https://gitlab.com/gitlab-org/gitter/webapp")

	session, err := s.Client.NewUploadPackSession(endpoint, s.EmptyAuth)
	c.Assert(err, IsNil)

	info, err := session.AdvertisedReferences()
	c.Assert(err, IsNil)
	c.Assert(info, NotNil)

	url := session.(*upSession).endpoint.String()
	c.Assert(url, Equals, "https://gitlab.com/gitlab-org/gitter/webapp.git")
}

func (s *UploadPackSuite) TestAdvertisedReferencesRedirectSchema(c *C) {
	endpoint, _ := transport.NewEndpoint("http://github.com/git-fixtures/basic")

	session, err := s.Client.NewUploadPackSession(endpoint, s.EmptyAuth)
	c.Assert(err, IsNil)

	info, err := session.AdvertisedReferences()
	c.Assert(err, IsNil)
	c.Assert(info, NotNil)

	url := session.(*upSession).endpoint.String()
	c.Assert(url, Equals, "https://github.com/git-fixtures/basic")
}

func (s *UploadPackSuite) TestAdvertisedReferencesContext(c *C) {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	endpoint, _ := transport.NewEndpoint("http://github.com/git-fixtures/basic")

	session, err := s.Client.NewUploadPackSession(endpoint, s.EmptyAuth)
	c.Assert(err, IsNil)

	info, err := session.AdvertisedReferencesContext(ctx)
	c.Assert(err, IsNil)
	c.Assert(info, NotNil)

	url := session.(*upSession).endpoint.String()
	c.Assert(url, Equals, "https://github.com/git-fixtures/basic")
}

func (s *UploadPackSuite) TestAdvertisedReferencesContextCanceled(c *C) {
	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	endpoint, _ := transport.NewEndpoint("http://github.com/git-fixtures/basic")

	session, err := s.Client.NewUploadPackSession(endpoint, s.EmptyAuth)
	c.Assert(err, IsNil)

	info, err := session.AdvertisedReferencesContext(ctx)
	c.Assert(err, DeepEquals, &url.Error{Op: "Get", URL: "http://github.com/git-fixtures/basic/info/refs?service=git-upload-pack", Err: context.Canceled})
	c.Assert(info, IsNil)
}

func (s *UploadPackSuite) TestUploadPackWithContextOnRead(c *C) {
	c.Skip("flaky tests, looks like sometimes the request body is cached, so doesn't fail on context cancel")
}
<|endoftext|>


<commit_before>// Copyright 2011 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 template (html/template) implements data-driven templates for
generating HTML output safe against code injection. It provides the
same interface as package text/template and should be used instead of
text/template whenever the output is HTML.

The documentation here focuses on the security features of the package.
For information about how to program the templates themselves, see the
documentation for text/template.

Introduction

This package wraps package text/template so you can share its template API
to parse and execute HTML templates safely.

  tmpl, err := template.New("name").Parse(...)
  // Error checking elided
  err = tmpl.Execute(out, data)

If successful, tmpl will now be injection-safe. Otherwise, err is an error
defined in the docs for ErrorCode.

HTML templates treat data values as plain text which should be encoded so they
can be safely embedded in an HTML document. The escaping is contextual, so
actions can appear within JavaScript, CSS, and URI contexts.

The security model used by this package assumes that template authors are
trusted, while Execute's data parameter is not. More details are
provided below.

Example

  import "text/template"
  ...
  t, err := template.New("foo").Parse(`{{define "T"}}Hello, {{.}}!{{end}}`)
  err = t.ExecuteTemplate(out, "T", "")

produces

  Hello, !

but the contextual autoescaping in html/template

  import "html/template"
  ...
  t, err := template.New("foo").Parse(`{{define "T"}}Hello, {{.}}!{{end}}`)
  err = t.ExecuteTemplate(out, "T", "")

produces safe, escaped HTML output

  Hello, <script>alert('you have been pwned')</script>!

Contexts

This package understands HTML, CSS, JavaScript, and URIs. It adds sanitizing
functions to each simple action pipeline, so given the excerpt

  {{.}}

At parse time each {{.}} is overwritten to add escaping functions as necessary.
In this case it becomes

  {{. | html}}

Errors

See the documentation of ErrorCode for details.

A fuller picture

The rest of this package comment may be skipped on first reading; it includes
details necessary to understand escaping contexts and error messages. Most users
will not need to understand these details.

Contexts

Assuming {{.}} is `O'Reilly: How are <i>you</i>?`, the table below shows
how {{.}} appears when used in the context to the left.

  Context                          {{.}} After
  {{.}}                            O'Reilly: How are <i>you</i>?
                  O'Reilly: How are you?
                  O'Reilly: How are %3ci%3eyou%3c/i%3e?
                O'Reilly%3a%20How%20are%3ci%3e...%3f
               O\x27Reilly: How are \x3ci\x3eyou...?
                 "O\x27Reilly: How are \x3ci\x3eyou...?"
       O\x27Reilly: How are \x3ci\x3eyou...\x3f

If used in an unsafe context, then the value might be filtered out:

  Context                          {{.}} After
                   #ZgotmplZ

since "O'Reilly:" is not an allowed protocol like "http:".

If {{.}} is the innocuous word, `left`, then it can appear more widely,

  Context                              {{.}} After
  {{.}}                                left
                      left
                       left
                      left
                  left
          left
               left
         left
    left
     left

Non-string values can be used in JavaScript contexts.
If {{.}} is

  struct{A,B string}{ "foo", "bar" }

in the escaped template

then the template output is

See package json to understand how non-string content is marshalled for
embedding in JavaScript contexts.

Typed Strings

By default, this package assumes that all pipelines produce a plain text string.
It adds escaping pipeline stages necessary to correctly and safely embed that
plain text string in the appropriate context.

When a data value is not plain text, you can make sure it is not over-escaped
by marking it with its type.

Types HTML, JS, URL, and others from content.go can carry safe content that is
exempted from escaping.

The template

  Hello, {{.}}!

can be invoked with

  tmpl.Execute(out, template.HTML(`<b>World</b>`))

to produce

  Hello, World!

instead of the

  Hello, <b>World<b>!

that would have been produced if {{.}} was a regular string.

Security Model

http://js-quasis-libraries-and-repl.googlecode.com/svn/trunk/safetemplate.html#problem_definition defines "safe" as used by this package.

This package assumes that template authors are trusted, that Execute's data
parameter is not, and seeks to preserve the properties below in the face
of untrusted data:

Structure Preservation Property:
"... when a template author writes an HTML tag in a safe templating language,
the browser will interpret the corresponding portion of the output as a tag
regardless of the values of untrusted data, and similarly for other structures
such as attribute boundaries and JS and CSS string boundaries."

Code Effect Property:
"... only code specified by the template author should run as a result of
injecting the template output into a page and all code specified by the
template author should run as a result of the same."

Least Surprise Property:
"A developer (or code reviewer) familiar with HTML, CSS, and JavaScript, who
knows that contextual autoescaping happens should be able to look at a {{.}}
and correctly infer what sanitization happens."
*/
package template
<commit_msg>html/template: update security model link<commit_after>// Copyright 2011 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 template (html/template) implements data-driven templates for
generating HTML output safe against code injection. It provides the
same interface as package text/template and should be used instead of
text/template whenever the output is HTML.

The documentation here focuses on the security features of the package.
For information about how to program the templates themselves, see the
documentation for text/template.

Introduction

This package wraps package text/template so you can share its template API
to parse and execute HTML templates safely.

  tmpl, err := template.New("name").Parse(...)
  // Error checking elided
  err = tmpl.Execute(out, data)

If successful, tmpl will now be injection-safe. Otherwise, err is an error
defined in the docs for ErrorCode.

HTML templates treat data values as plain text which should be encoded so they
can be safely embedded in an HTML document. The escaping is contextual, so
actions can appear within JavaScript, CSS, and URI contexts.

The security model used by this package assumes that template authors are
trusted, while Execute's data parameter is not. More details are
provided below.

Example

  import "text/template"
  ...
  t, err := template.New("foo").Parse(`{{define "T"}}Hello, {{.}}!{{end}}`)
  err = t.ExecuteTemplate(out, "T", "")

produces

  Hello, !

but the contextual autoescaping in html/template

  import "html/template"
  ...
  t, err := template.New("foo").Parse(`{{define "T"}}Hello, {{.}}!{{end}}`)
  err = t.ExecuteTemplate(out, "T", "")

produces safe, escaped HTML output

  Hello, <script>alert('you have been pwned')</script>!

Contexts

This package understands HTML, CSS, JavaScript, and URIs. It adds sanitizing
functions to each simple action pipeline, so given the excerpt

  {{.}}

At parse time each {{.}} is overwritten to add escaping functions as necessary.
In this case it becomes

  {{. | html}}

Errors

See the documentation of ErrorCode for details.

A fuller picture

The rest of this package comment may be skipped on first reading; it includes
details necessary to understand escaping contexts and error messages. Most users
will not need to understand these details.

Contexts

Assuming {{.}} is `O'Reilly: How are <i>you</i>?`, the table below shows
how {{.}} appears when used in the context to the left.

  Context                          {{.}} After
  {{.}}                            O'Reilly: How are <i>you</i>?
                  O'Reilly: How are you?
                  O'Reilly: How are %3ci%3eyou%3c/i%3e?
                O'Reilly%3a%20How%20are%3ci%3e...%3f
               O\x27Reilly: How are \x3ci\x3eyou...?
                 "O\x27Reilly: How are \x3ci\x3eyou...?"
       O\x27Reilly: How are \x3ci\x3eyou...\x3f

If used in an unsafe context, then the value might be filtered out:

  Context                          {{.}} After
                   #ZgotmplZ

since "O'Reilly:" is not an allowed protocol like "http:".

If {{.}} is the innocuous word, `left`, then it can appear more widely,

  Context                              {{.}} After
  {{.}}                                left
                      left
                       left
                      left
                  left
          left
               left
         left
    left
     left

Non-string values can be used in JavaScript contexts.
If {{.}} is

  struct{A,B string}{ "foo", "bar" }

in the escaped template

then the template output is

See package json to understand how non-string content is marshalled for
embedding in JavaScript contexts.

Typed Strings

By default, this package assumes that all pipelines produce a plain text string.
It adds escaping pipeline stages necessary to correctly and safely embed that
plain text string in the appropriate context.

When a data value is not plain text, you can make sure it is not over-escaped
by marking it with its type.

Types HTML, JS, URL, and others from content.go can carry safe content that is
exempted from escaping.

The template

  Hello, {{.}}!

can be invoked with

  tmpl.Execute(out, template.HTML(`<b>World</b>`))

to produce

  Hello, World!

instead of the

  Hello, <b>World<b>!

that would have been produced if {{.}} was a regular string.

Security Model

https://rawgit.com/mikesamuel/sanitized-jquery-templates/trunk/safetemplate.html#problem_definition defines "safe" as used by this package.

This package assumes that template authors are trusted, that Execute's data
parameter is not, and seeks to preserve the properties below in the face
of untrusted data:

Structure Preservation Property:
"... when a template author writes an HTML tag in a safe templating language,
the browser will interpret the corresponding portion of the output as a tag
regardless of the values of untrusted data, and similarly for other structures
such as attribute boundaries and JS and CSS string boundaries."

Code Effect Property:
"... only code specified by the template author should run as a result of
injecting the template output into a page and all code specified by the
template author should run as a result of the same."

Least Surprise Property:
"A developer (or code reviewer) familiar with HTML, CSS, and JavaScript, who
knows that contextual autoescaping happens should be able to look at a {{.}}
and correctly infer what sanitization happens."
*/
package template
<|endoftext|>


<commit_before>// Copyright 2012 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 main

import (
	"bytes"
	"github.com/globocom/commandmocker"
	"github.com/globocom/tsuru/api/bind"
	"github.com/globocom/tsuru/app"
	"github.com/globocom/tsuru/db"
	"github.com/globocom/tsuru/log"
	"github.com/globocom/tsuru/queue"
	"labix.org/v2/mgo/bson"
	. "launchpad.net/gocheck"
	stdlog "log"
	"strings"
	"time"
)

func (s *S) TestHandleMessages(c *C) {
	handler := MessageHandler{}
	err := handler.start()
	c.Assert(err, IsNil)
	defer handler.stop()
	a := app.App{
		Name: "nemesis",
		Units: []app.Unit{
			{
				Name:    "i-00800",
				State:   "started",
				Machine: 19,
			},
		},
		Env: map[string]bind.EnvVar{
			"http_proxy": {
				Name:   "http_proxy",
				Value:  "http://myproxy.com:3128/",
				Public: true,
			},
		},
		State: "started",
	}
	err = db.Session.Apps().Insert(a)
	c.Assert(err, IsNil)
	defer db.Session.Apps().Remove(bson.M{"name": a.Name})
	tmpdir, err := commandmocker.Add("juju", "$*")
	c.Assert(err, IsNil)
	defer commandmocker.Remove(tmpdir)
	messages, _, err := queue.Dial(handler.server.Addr())
	c.Assert(err, IsNil)
	messages <- queue.Message{Action: app.RegenerateApprc, Args: []string{a.Name}}
	time.Sleep(1e9)
	c.Assert(commandmocker.Ran(tmpdir), Equals, true)
	output := strings.Replace(commandmocker.Output(tmpdir), "\n", " ", -1)
	outputRegexp := `^.*19 cat > /home/application/apprc <<END # generated by tsuru.*`
	outputRegexp += `export http_proxy="http://myproxy.com:3128/" END $`
	c.Assert(output, Matches, outputRegexp)
}

func (s *S) TestHandleMessageErrors(c *C) {
	var data = []struct {
		action      string
		appName     string
		expectedLog string
		visits      int
	}{
		{
			action:      "unknown-action",
			appName:     "does not matter",
			expectedLog: `Error handling "unknown-action": invalid action.`,
		},
		{
			action:  app.RegenerateApprc,
			appName: "nemesis",
			expectedLog: `Error handling "regenerate-apprc" for the app "nemesis":` +
				` The status of the app should be "started", but it is "pending".`,
		},
		{
			action:      app.RegenerateApprc,
			appName:     "unknown-app",
			expectedLog: `Error handling "regenerate-apprc": app "unknown-app" does not exist.`,
		},
		{
			action:      app.RegenerateApprc,
			expectedLog: `Error handling "regenerate-apprc": this action requires at least 1 argument.`,
		},
		{
			action:      "does not matter",
			appName:     "does not matter",
			expectedLog: `Error handling "does not matter": this message has been visited more than 50 times.`,
			visits:      MaxVisits,
		},
		{
			action:  app.RegenerateApprc,
			appName: "marathon",
			expectedLog: `Error handling "regenerate-apprc" for the app "marathon":` +
				` the app is in "error" state.`,
		},
		{
			action:  app.RegenerateApprc,
			appName: "territories",
			expectedLog: `Error handling "regenerate-apprc" for the app "territories":` +
				` the app is down.`,
		},
	}
	var buf bytes.Buffer
	a := app.App{Name: "nemesis", State: "pending"}
	err := db.Session.Apps().Insert(a)
	c.Assert(err, IsNil)
	defer db.Session.Apps().Remove(bson.M{"name": a.Name})
	a = app.App{Name: "marathon", State: "error"}
	err = db.Session.Apps().Insert(a)
	c.Assert(err, IsNil)
	defer db.Session.Apps().Remove(bson.M{"name": a.Name})
	a = app.App{Name: "territories", State: "down"}
	err = db.Session.Apps().Insert(a)
	c.Assert(err, IsNil)
	defer db.Session.Apps().Remove(bson.M{"name": a.Name})
	log.SetLogger(stdlog.New(&buf, "", 0))
	handler := MessageHandler{}
	handler.start()
	handler.closed = 1
	defer handler.stop()
	for _, d := range data {
		message := queue.Message{
			Action: d.action,
			Visits: d.visits,
		}
		if d.appName != "" {
			message.Args = append(message.Args, d.appName)
		}
		handler.handle(message)
	}
	content := buf.String()
	lines := strings.Split(content, "\n")
	for i, d := range data {
		var found bool
		for j := i; j < len(lines); j++ {
			if lines[j] == d.expectedLog {
				found = true
				break
			}
		}
		if !found {
			c.Errorf("\nWant: %q.\nGot:\n%s", d.expectedLog, content)
		}
	}
}
<commit_msg>collector: fix tests to use new app package<commit_after>// Copyright 2012 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 main

import (
	"bytes"
	"github.com/globocom/tsuru/api/bind"
	"github.com/globocom/tsuru/app"
	"github.com/globocom/tsuru/db"
	"github.com/globocom/tsuru/log"
	"github.com/globocom/tsuru/provision"
	"github.com/globocom/tsuru/queue"
	"labix.org/v2/mgo/bson"
	. "launchpad.net/gocheck"
	stdlog "log"
	"strings"
	"time"
)

func (s *S) TestHandleMessages(c *C) {
	s.provisioner.PrepareOutput([]byte("exported"))
	handler := MessageHandler{}
	err := handler.start()
	c.Assert(err, IsNil)
	defer handler.stop()
	a := app.App{
		Name: "nemesis",
		Units: []app.Unit{
			{
				Name:    "i-00800",
				State:   "started",
				Machine: 19,
			},
		},
		Env: map[string]bind.EnvVar{
			"http_proxy": {
				Name:   "http_proxy",
				Value:  "http://myproxy.com:3128/",
				Public: true,
			},
		},
		State: provision.StatusStarted,
	}
	err = db.Session.Apps().Insert(a)
	c.Assert(err, IsNil)
	defer db.Session.Apps().Remove(bson.M{"name": a.Name})
	messages, _, err := queue.Dial(handler.server.Addr())
	c.Assert(err, IsNil)
	messages <- queue.Message{Action: app.RegenerateApprc, Args: []string{a.Name}}
	time.Sleep(1e9)
	cmds := s.provisioner.GetCmds("", &a)
	c.Assert(cmds, HasLen, 1)
	output := strings.Replace(cmds[0].Cmd, "\n", " ", -1)
	outputRegexp := `^cat > /home/application/apprc <<END # generated by tsuru.*`
	outputRegexp += `export http_proxy="http://myproxy.com:3128/" END $`
	c.Assert(output, Matches, outputRegexp)
}

func (s *S) TestHandleMessageErrors(c *C) {
	var data = []struct {
		action      string
		appName     string
		expectedLog string
		visits      int
	}{
		{
			action:      "unknown-action",
			appName:     "does not matter",
			expectedLog: `Error handling "unknown-action": invalid action.`,
		},
		{
			action:  app.RegenerateApprc,
			appName: "nemesis",
			expectedLog: `Error handling "regenerate-apprc" for the app "nemesis":` +
				` The status of the app should be "started", but it is "pending".`,
		},
		{
			action:      app.RegenerateApprc,
			appName:     "unknown-app",
			expectedLog: `Error handling "regenerate-apprc": app "unknown-app" does not exist.`,
		},
		{
			action:      app.RegenerateApprc,
			expectedLog: `Error handling "regenerate-apprc": this action requires at least 1 argument.`,
		},
		{
			action:      "does not matter",
			appName:     "does not matter",
			expectedLog: `Error handling "does not matter": this message has been visited more than 50 times.`,
			visits:      MaxVisits,
		},
		{
			action:  app.RegenerateApprc,
			appName: "marathon",
			expectedLog: `Error handling "regenerate-apprc" for the app "marathon":` +
				` the app is in "error" state.`,
		},
		{
			action:  app.RegenerateApprc,
			appName: "territories",
			expectedLog: `Error handling "regenerate-apprc" for the app "territories":` +
				` the app is down.`,
		},
	}
	var buf bytes.Buffer
	a := app.App{Name: "nemesis", State: "pending"}
	err := db.Session.Apps().Insert(a)
	c.Assert(err, IsNil)
	defer db.Session.Apps().Remove(bson.M{"name": a.Name})
	a = app.App{Name: "marathon", State: "error"}
	err = db.Session.Apps().Insert(a)
	c.Assert(err, IsNil)
	defer db.Session.Apps().Remove(bson.M{"name": a.Name})
	a = app.App{Name: "territories", State: "down"}
	err = db.Session.Apps().Insert(a)
	c.Assert(err, IsNil)
	defer db.Session.Apps().Remove(bson.M{"name": a.Name})
	log.SetLogger(stdlog.New(&buf, "", 0))
	handler := MessageHandler{}
	handler.start()
	handler.closed = 1
	defer handler.stop()
	for _, d := range data {
		message := queue.Message{
			Action: d.action,
			Visits: d.visits,
		}
		if d.appName != "" {
			message.Args = append(message.Args, d.appName)
		}
		handler.handle(message)
	}
	content := buf.String()
	lines := strings.Split(content, "\n")
	for i, d := range data {
		var found bool
		for j := i; j < len(lines); j++ {
			if lines[j] == d.expectedLog {
				found = true
				break
			}
		}
		if !found {
			c.Errorf("\nWant: %q.\nGot:\n%s", d.expectedLog, content)
		}
	}
}
<|endoftext|>


<commit_before>package command

import (
	"fmt"
	"testing"
)

func TestColor2String(t *testing.T) {
	cases := []struct {
		color Color
		str   string
	}{
		{Color{0, 0, 0}, "#000000"},
		{Color{255, 0, 0}, "#ff0000"},
		{Color{0, 128, 0}, "#008000"},
		{Color{0, 0, 15}, "#00000f"},
	}

	for _, c := range cases {
		got := fmt.Sprintf("%s", c.color)
		if got != c.str {
			t.Errorf("Color{%d, %d, %d} - expected %q, got %q", c.color.r, c.color.g, c.color.b, c.str, got)
		}
	}
}
<commit_msg>Add command test<commit_after>package command

import (
	"encoding/json"
	"fmt"
	"testing"
)

func TestColor2String(t *testing.T) {
	cases := []struct {
		color Color
		str   string
	}{
		{Color{0, 0, 0}, "#000000"},
		{Color{255, 0, 0}, "#ff0000"},
		{Color{0, 128, 0}, "#008000"},
		{Color{0, 0, 15}, "#00000f"},
	}

	for _, c := range cases {
		got := fmt.Sprintf("%s", c.color)
		if got != c.str {
			t.Errorf("Color{%d, %d, %d} - expected %q, got %q", c.color.r, c.color.g, c.color.b, c.str, got)
		}
	}
}

func TestColor2Json(t *testing.T) {
	cases := []struct {
		color Color
		str   string
	}{
		{Color{0, 0, 0}, `"#000000"`},
		{Color{255, 0, 0}, `"#ff0000"`},
		{Color{0, 128, 0}, `"#008000"`},
		{Color{0, 0, 15}, `"#00000f"`},
	}

	for _, c := range cases {
		buf, e := json.Marshal(c.color)
		if e != nil {
			t.Errorf("Error occured while converting to json : %q", e)
		}
		got := string(buf)
		if got != c.str {
			t.Errorf("json.Marshal(Color{%d, %d, %d}) - expected %q, got %q", c.color.r, c.color.g, c.color.b, c.str, got)
		}
	}
}
<|endoftext|>


<commit_before>package command

import (
	"bytes"
	"fmt"
	"sort"
	"strings"

	"github.com/hashicorp/terraform/terraform"
	"github.com/mitchellh/colorstring"
)

// FormatStateOpts are the options for formatting a state.
type FormatStateOpts struct {
	// State is the state to format. This is required.
	State *terraform.State

	// Color is the colorizer. This is optional.
	Color *colorstring.Colorize

	// ModuleDepth is the depth of the modules to expand. By default this
	// is zero which will not expand modules at all.
	ModuleDepth int
}

// FormatState takes a state and returns a string
func FormatState(opts *FormatStateOpts) string {
	if opts.Color == nil {
		panic("colorize not given")
	}

	s := opts.State
	if len(s.Modules) == 0 {
		return "The state file is empty. No resources are represented."
	}

	var buf bytes.Buffer
	buf.WriteString("[reset]")

	// Format all the modules
	for _, m := range s.Modules {
		if len(m.Path)-1 <= opts.ModuleDepth || opts.ModuleDepth == -1 {
			formatStateModuleExpand(&buf, m, opts)
		} else {
			formatStateModuleSingle(&buf, m, opts)
		}
	}

	// Write the outputs for the root module
	m := s.RootModule()
	if len(m.Outputs) > 0 {
		buf.WriteString("\nOutputs:\n\n")

		// Sort the outputs
		ks := make([]string, 0, len(m.Outputs))
		for k, _ := range m.Outputs {
			ks = append(ks, k)
		}
		sort.Strings(ks)

		// Output each output k/v pair
		for _, k := range ks {
			v := m.Outputs[k]
			buf.WriteString(fmt.Sprintf("%s = %s\n", k, v))
		}
	}

	return opts.Color.Color(strings.TrimSpace(buf.String()))
}

func formatStateModuleExpand(
	buf *bytes.Buffer, m *terraform.ModuleState, opts *FormatStateOpts) {
	var moduleName string
	if !m.IsRoot() {
		moduleName = fmt.Sprintf("module.%s", strings.Join(m.Path[1:], "."))
	}

	// First get the names of all the resources so we can show them
	// in alphabetical order.
	names := make([]string, 0, len(m.Resources))
	for name, _ := range m.Resources {
		names = append(names, name)
	}
	sort.Strings(names)

	// Go through each resource and begin building up the output.
	for _, k := range names {
		name := k
		if moduleName != "" {
			name = moduleName + "." + name
		}

		rs := m.Resources[k]
		is := rs.Primary
		id := is.ID
		if id == "" {
			id = "<not created>"
		}

		taintStr := ""
		if len(rs.Tainted) > 0 {
			taintStr = " (tainted)"
		}

		buf.WriteString(fmt.Sprintf("%s:%s\n", name, taintStr))
		buf.WriteString(fmt.Sprintf("  id = %s\n", id))

		// Sort the attributes
		attrKeys := make([]string, 0, len(is.Attributes))
		for ak, _ := range is.Attributes {
			// Skip the id attribute since we just show the id directly
			if ak == "id" {
				continue
			}

			attrKeys = append(attrKeys, ak)
		}
		sort.Strings(attrKeys)

		// Output each attribute
		for _, ak := range attrKeys {
			av := is.Attributes[ak]
			buf.WriteString(fmt.Sprintf("  %s = %s\n", ak, av))
		}
	}

	buf.WriteString("[reset]\n")
}

func formatStateModuleSingle(
	buf *bytes.Buffer, m *terraform.ModuleState, opts *FormatStateOpts) {
	// Header with the module name
	buf.WriteString(fmt.Sprintf("module.%s\n", strings.Join(m.Path[1:], ".")))

	// Now just write how many resources are in here.
	buf.WriteString(fmt.Sprintf("  %d resource(s)\n", len(m.Resources)))
}
<commit_msg>command: nil-check for instance state when showing state<commit_after>package command

import (
	"bytes"
	"fmt"
	"sort"
	"strings"

	"github.com/hashicorp/terraform/terraform"
	"github.com/mitchellh/colorstring"
)

// FormatStateOpts are the options for formatting a state.
type FormatStateOpts struct {
	// State is the state to format. This is required.
	State *terraform.State

	// Color is the colorizer. This is optional.
	Color *colorstring.Colorize

	// ModuleDepth is the depth of the modules to expand. By default this
	// is zero which will not expand modules at all.
	ModuleDepth int
}

// FormatState takes a state and returns a string
func FormatState(opts *FormatStateOpts) string {
	if opts.Color == nil {
		panic("colorize not given")
	}

	s := opts.State
	if len(s.Modules) == 0 {
		return "The state file is empty. No resources are represented."
	}

	var buf bytes.Buffer
	buf.WriteString("[reset]")

	// Format all the modules
	for _, m := range s.Modules {
		if len(m.Path)-1 <= opts.ModuleDepth || opts.ModuleDepth == -1 {
			formatStateModuleExpand(&buf, m, opts)
		} else {
			formatStateModuleSingle(&buf, m, opts)
		}
	}

	// Write the outputs for the root module
	m := s.RootModule()
	if len(m.Outputs) > 0 {
		buf.WriteString("\nOutputs:\n\n")

		// Sort the outputs
		ks := make([]string, 0, len(m.Outputs))
		for k, _ := range m.Outputs {
			ks = append(ks, k)
		}
		sort.Strings(ks)

		// Output each output k/v pair
		for _, k := range ks {
			v := m.Outputs[k]
			buf.WriteString(fmt.Sprintf("%s = %s\n", k, v))
		}
	}

	return opts.Color.Color(strings.TrimSpace(buf.String()))
}

func formatStateModuleExpand(
	buf *bytes.Buffer, m *terraform.ModuleState, opts *FormatStateOpts) {
	var moduleName string
	if !m.IsRoot() {
		moduleName = fmt.Sprintf("module.%s", strings.Join(m.Path[1:], "."))
	}

	// First get the names of all the resources so we can show them
	// in alphabetical order.
	names := make([]string, 0, len(m.Resources))
	for name, _ := range m.Resources {
		names = append(names, name)
	}
	sort.Strings(names)

	// Go through each resource and begin building up the output.
	for _, k := range names {
		name := k
		if moduleName != "" {
			name = moduleName + "." + name
		}

		rs := m.Resources[k]
		is := rs.Primary
		var id string
		if is != nil {
			id = is.ID
		}
		if id == "" {
			id = "<not created>"
		}

		taintStr := ""
		if len(rs.Tainted) > 0 {
			taintStr = " (tainted)"
		}

		buf.WriteString(fmt.Sprintf("%s:%s\n", name, taintStr))
		buf.WriteString(fmt.Sprintf("  id = %s\n", id))

		if is != nil {
			// Sort the attributes
			attrKeys := make([]string, 0, len(is.Attributes))
			for ak, _ := range is.Attributes {
				// Skip the id attribute since we just show the id directly
				if ak == "id" {
					continue
				}

				attrKeys = append(attrKeys, ak)
			}
			sort.Strings(attrKeys)

			// Output each attribute
			for _, ak := range attrKeys {
				av := is.Attributes[ak]
				buf.WriteString(fmt.Sprintf("  %s = %s\n", ak, av))
			}
		}
	}

	buf.WriteString("[reset]\n")
}

func formatStateModuleSingle(
	buf *bytes.Buffer, m *terraform.ModuleState, opts *FormatStateOpts) {
	// Header with the module name
	buf.WriteString(fmt.Sprintf("module.%s\n", strings.Join(m.Path[1:], ".")))

	// Now just write how many resources are in here.
	buf.WriteString(fmt.Sprintf("  %d resource(s)\n", len(m.Resources)))
}
<|endoftext|>


<commit_before>package command

import (
	"log"

	"github.com/codegangsta/cli"
)

var PushFlags = []cli.Flag{
	cli.StringFlag{
		Name:  "base-dir",
		Usage: "base directory of target files",
	},
}

func PushCommand(c *cli.Context) {
	base_dir := c.String("base-dir")
	if base_dir == "" {
		log.Fatalf("Error missing flag 'base-dir'")
	}
}
<commit_msg>Add push command<commit_after>package command

import (
	"github.com/codegangsta/cli"
	"log"

	. "github.com/foostan/fileconsul/fileconsul"
	"path/filepath"
)

var PushFlags = []cli.Flag{
	cli.StringFlag{
		Name:  "addr",
		Value: "localhost:8500",
		Usage: "consul HTTP API address with port",
	},
	cli.StringFlag{
		Name:  "dc",
		Value: "dc1",
		Usage: "consul datacenter, uses local if blank",
	},
	cli.StringFlag{
		Name:  "prefix",
		Value: "fileconsul",
		Usage: "reading file status from Consul's K/V store with the given prefix",
	},
	cli.StringFlag{
		Name:  "base-dir",
		Value: ".",
		Usage: "base directory of target files",
	},
}

func PushCommand(c *cli.Context) {
	addr := c.String("addr")
	dc := c.String("dc")
	prefix := c.String("prefix")
	baseDir := c.String("base-dir")

	client, err := NewClient(&ClientConfig{
		ConsulAddr: addr,
		ConsulDC:   dc,
	})
	if err != nil {
		log.Fatal(err)
	}

	localFhs, err := LocalFileHashs(baseDir)
	if err != nil {
		log.Fatal(err)
	}

	remoteFhs, err := RemoteFileHashs(client, prefix)
	if err != nil {
		log.Fatal(err)
	}

	addFhs, delFhs, modFhs := DiffFileHashs(localFhs, remoteFhs)
	for _, fh := range addFhs {
		println("push new file:\t" + fh.Path)
		err = client.PutKV(filepath.Join(prefix, fh.Path), fh.Hash)
		if err != nil {
			log.Fatal(err)
		}
	}
	for _, fh := range delFhs {
		println("delete file:\t" + fh.Path)
		err = client.DeleteKV(filepath.Join(prefix, fh.Path))
		if err != nil {
			log.Fatal(err)
		}
	}
	for _, fh := range modFhs {
		println("modify file:\t" + fh.Path)
		err = client.PutKV(filepath.Join(prefix, fh.Path), fh.Hash)
		if err != nil {
			log.Fatal(err)
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"github.com/ryanbressler/CloudForest"
	"log"
	"os"
	"runtime"
	"strings"
)

func main() {
	fm := flag.String("fm", "featurematrix.afm", "AFM formated feature matrix to use.")
	rf := flag.String("rfpred", "rface.sf", "A predictor forest.")
	outf := flag.String("leaves", "leaves.tsv", "a case by case sparse matrix of leaf co-occurrence in tsv format")
	boutf := flag.String("branches", "", "a case by feature sparse matrix of leaf co-occurrence in tsv format")
	var threads int
	flag.IntVar(&threads, "threads", 1, "Parse seperate forests in n seperate threads.")

	flag.Parse()

	//Parse Data
	data, err := CloudForest.LoadAFM(*fm)
	if err != nil {
		log.Fatal(err)
	}

	log.Print("Data file ", len(data.Data), " by ", data.Data[0].Length())

	counts := new(CloudForest.SparseCounter)
	var caseFeatureCounts *CloudForest.SparseCounter
	if *boutf != "" {
		caseFeatureCounts = new(CloudForest.SparseCounter)
	}

	files := strings.Split(*rf, ",")

	runtime.GOMAXPROCS(threads)

	fileChan := make(chan string, 0)
	doneChan := make(chan int, 0)

	go func() {
		for _, fn := range files {
			fileChan <- fn
		}
	}()

	for i := 0; i < threads; i++ {

		go func() {
			for {
				fn := <-fileChan

				forestfile, err := os.Open(fn) // For read access.
				if err != nil {
					log.Fatal(err)
				}
				defer forestfile.Close()
				forestreader := CloudForest.NewForestReader(forestfile)
				forest, err := forestreader.ReadForest()
				if err != nil {
					log.Fatal(err)
				}
				log.Print("Forest has ", len(forest.Trees), " trees ")

				for i := 0; i < len(forest.Trees); i++ {
					leaves := forest.Trees[i].GetLeaves(data, caseFeatureCounts)
					for _, leaf := range leaves {
						for j := 0; j < len(leaf.Cases); j++ {
							for k := 0; k < len(leaf.Cases); k++ {

								counts.Add(leaf.Cases[j], leaf.Cases[k], 1)

							}
						}
					}

				}
				doneChan <- 1
			}
		}()

	}

	for i := 0; i < len(files); i++ {
		<-doneChan
	}

	log.Print("Outputting Case Case  Co-Occurrence Counts")
	outfile, err := os.Create(*outf) // For read access.
	if err != nil {
		log.Fatal(err)
	}
	defer outfile.Close()
	counts.WriteTsv(outfile)

	if *boutf != "" {
		log.Print("Outputting Case Feature Co-Occurrence Counts")
		boutfile, err := os.Create(*boutf) // For read access.
		if err != nil {
			log.Fatal(err)
		}
		defer boutfile.Close()
		caseFeatureCounts.WriteTsv(boutfile)
	}
}
<commit_msg>dump slit values from leafcount<commit_after>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"github.com/ryanbressler/CloudForest"
	"log"
	"os"
	"runtime"
	"strings"
)

func main() {
	fm := flag.String("fm", "featurematrix.afm", "AFM formated feature matrix to use.")
	rf := flag.String("rfpred", "rface.sf", "A predictor forest.")
	outf := flag.String("leaves", "leaves.tsv", "a case by case sparse matrix of leaf co-occurrence in tsv format")
	boutf := flag.String("branches", "", "a case by feature sparse matrix of leaf co-occurrence in tsv format")
	soutf := flag.String("splits", "", "a file to write a json record of splite per feature")
	var threads int
	flag.IntVar(&threads, "threads", 1, "Parse seperate forests in n seperate threads.")

	flag.Parse()

	splits := make(map[string][]string)

	//Parse Data
	data, err := CloudForest.LoadAFM(*fm)
	if err != nil {
		log.Fatal(err)
	}

	log.Print("Data file ", len(data.Data), " by ", data.Data[0].Length())

	counts := new(CloudForest.SparseCounter)
	var caseFeatureCounts *CloudForest.SparseCounter
	if *boutf != "" {
		caseFeatureCounts = new(CloudForest.SparseCounter)
	}

	files := strings.Split(*rf, ",")

	runtime.GOMAXPROCS(threads)

	fileChan := make(chan string, 0)
	doneChan := make(chan int, 0)

	go func() {
		for _, fn := range files {
			fileChan <- fn
		}
	}()

	for i := 0; i < threads; i++ {

		go func() {
			for {
				fn := <-fileChan

				forestfile, err := os.Open(fn) // For read access.
				if err != nil {
					log.Fatal(err)
				}
				defer forestfile.Close()
				forestreader := CloudForest.NewForestReader(forestfile)
				forest, err := forestreader.ReadForest()
				if err != nil {
					log.Fatal(err)
				}
				log.Print("Forest has ", len(forest.Trees), " trees ")

				for i := 0; i < len(forest.Trees); i++ {
					fmt.Print(".")
					leaves := forest.Trees[i].GetLeaves(data, caseFeatureCounts)
					for _, leaf := range leaves {
						for j := 0; j < len(leaf.Cases); j++ {
							for k := 0; k < len(leaf.Cases); k++ {

								counts.Add(leaf.Cases[j], leaf.Cases[k], 1)

							}
						}
					}

					if *soutf != "" {
						forest.Trees[i].Root.Climb(func(n *CloudForest.Node) {
							if n.Splitter != nil {
								name := n.Splitter.Feature
								_, ok := splits[name]
								if !ok {
									splits[name] = make([]string, 0, 10)
								}
								split := ""
								switch n.Splitter.Numerical {
								case true:
									split = fmt.Sprintf("%v", n.Splitter.Value)
								case false:
									keys := make([]string, 0, len(n.Splitter.Left))
									for k := range n.Splitter.Left {
										keys = append(keys, k)
									}
									split = strings.Join(keys, ",")
								}
								splits[name] = append(splits[name], split)
							}
						})
					}

				}
				doneChan <- 1
			}
		}()

	}

	for i := 0; i < len(files); i++ {
		<-doneChan
	}

	log.Print("Outputting Case Case  Co-Occurrence Counts")
	outfile, err := os.Create(*outf)
	if err != nil {
		log.Fatal(err)
	}
	defer outfile.Close()
	counts.WriteTsv(outfile)

	if *boutf != "" {
		log.Print("Outputting Case Feature Co-Occurrence Counts")
		boutfile, err := os.Create(*boutf)
		if err != nil {
			log.Fatal(err)
		}
		defer boutfile.Close()
		caseFeatureCounts.WriteTsv(boutfile)
	}

	if *soutf != "" {
		soutfile, err := os.Create(*soutf)
		if err != nil {
			log.Fatal(err)
		}
		defer soutfile.Close()
		encoder := json.NewEncoder(soutfile)
		encoder.Encode(splits)
	}
}
<|endoftext|>


<commit_before>// Copyright (c) 2012, Suryandaru Triandana <syndtr@gmail.com>
// All rights reserved.
//
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

// Package memdb provides in-memory key/value database implementation.
package memdb

import (
	"math/rand"
	"sync"

	"github.com/conformal/goleveldb/leveldb/comparer"
	"github.com/conformal/goleveldb/leveldb/iterator"
	"github.com/conformal/goleveldb/leveldb/util"
)

var (
	ErrNotFound = util.ErrNotFound
)

const tMaxHeight = 12

type dbIter struct {
	util.BasicReleaser
	p          *DB
	slice      *util.Range
	node       int
	forward    bool
	key, value []byte
}

func (i *dbIter) fill(checkStart, checkLimit bool) bool {
	if i.node != 0 {
		n := i.p.nodeData[i.node]
		m := n + i.p.nodeData[i.node+nKey]
		i.key = i.p.kvData[n:m]
		if i.slice != nil {
			switch {
			case checkLimit && i.slice.Limit != nil && i.p.cmp.Compare(i.key, i.slice.Limit) >= 0:
				fallthrough
			case checkStart && i.slice.Start != nil && i.p.cmp.Compare(i.key, i.slice.Start) < 0:
				i.node = 0
				goto bail
			}
		}
		i.value = i.p.kvData[m : m+i.p.nodeData[i.node+nVal]]
		return true
	}
bail:
	i.key = nil
	i.value = nil
	return false
}

func (i *dbIter) Valid() bool {
	return i.node != 0
}

func (i *dbIter) First() bool {
	i.forward = true
	i.p.mu.RLock()
	defer i.p.mu.RUnlock()
	if i.slice != nil && i.slice.Start != nil {
		i.node, _ = i.p.findGE(i.slice.Start, false)
	} else {
		i.node = i.p.nodeData[nNext]
	}
	return i.fill(false, true)
}

func (i *dbIter) Last() bool {
	i.forward = false
	i.p.mu.RLock()
	defer i.p.mu.RUnlock()
	if i.slice != nil && i.slice.Limit != nil {
		i.node = i.p.findLT(i.slice.Limit)
	} else {
		i.node = i.p.findLast()
	}
	return i.fill(true, false)
}

func (i *dbIter) Seek(key []byte) bool {
	i.forward = true
	i.p.mu.RLock()
	defer i.p.mu.RUnlock()
	if i.slice != nil && i.slice.Start != nil && i.p.cmp.Compare(key, i.slice.Start) < 0 {
		key = i.slice.Start
	}
	i.node, _ = i.p.findGE(key, false)
	return i.fill(false, true)
}

func (i *dbIter) Next() bool {
	if i.node == 0 {
		if !i.forward {
			return i.First()
		}
		return false
	}
	i.forward = true
	i.p.mu.RLock()
	defer i.p.mu.RUnlock()
	i.node = i.p.nodeData[i.node+nNext]
	return i.fill(false, true)
}

func (i *dbIter) Prev() bool {
	if i.node == 0 {
		if i.forward {
			return i.Last()
		}
		return false
	}
	i.forward = false
	i.p.mu.RLock()
	defer i.p.mu.RUnlock()
	i.node = i.p.findLT(i.key)
	return i.fill(true, false)
}

func (i *dbIter) Key() []byte {
	return i.key
}

func (i *dbIter) Value() []byte {
	return i.value
}

func (i *dbIter) Error() error { return nil }

const (
	nKV = iota
	nKey
	nVal
	nHeight
	nNext
)

// DB is an in-memory key/value database.
type DB struct {
	cmp comparer.BasicComparer
	rnd *rand.Rand

	mu     sync.RWMutex
	kvData []byte
	// Node data:
	// [0]         : KV offset
	// [1]         : Key length
	// [2]         : Value length
	// [3]         : Height
	// [3..height] : Next nodes
	nodeData  []int
	prevNode  [tMaxHeight]int
	maxHeight int
	n         int
	kvSize    int
}

func (p *DB) randHeight() (h int) {
	const branching = 4
	h = 1
	for h < tMaxHeight && p.rnd.Int()%branching == 0 {
		h++
	}
	return
}

func (p *DB) findGE(key []byte, prev bool) (int, bool) {
	node := 0
	h := p.maxHeight - 1
	for {
		next := p.nodeData[node+nNext+h]
		cmp := 1
		if next != 0 {
			o := p.nodeData[next]
			cmp = p.cmp.Compare(p.kvData[o:o+p.nodeData[next+nKey]], key)
		}
		if cmp < 0 {
			// Keep searching in this list
			node = next
		} else {
			if prev {
				p.prevNode[h] = node
			} else if cmp == 0 {
				return next, true
			}
			if h == 0 {
				return next, cmp == 0
			}
			h--
		}
	}
}

func (p *DB) findLT(key []byte) int {
	node := 0
	h := p.maxHeight - 1
	for {
		next := p.nodeData[node+nNext+h]
		o := p.nodeData[next]
		if next == 0 || p.cmp.Compare(p.kvData[o:o+p.nodeData[next+nKey]], key) >= 0 {
			if h == 0 {
				break
			}
			h--
		} else {
			node = next
		}
	}
	return node
}

func (p *DB) findLast() int {
	node := 0
	h := p.maxHeight - 1
	for {
		next := p.nodeData[node+nNext+h]
		if next == 0 {
			if h == 0 {
				break
			}
			h--
		} else {
			node = next
		}
	}
	return node
}

// Put sets the value for the given key. It overwrites any previous value
// for that key; a DB is not a multi-map.
//
// It is safe to modify the contents of the arguments after Put returns.
func (p *DB) Put(key []byte, value []byte) error {
	p.mu.Lock()
	defer p.mu.Unlock()

	if node, exact := p.findGE(key, true); exact {
		kvOffset := len(p.kvData)
		p.kvData = append(p.kvData, key...)
		p.kvData = append(p.kvData, value...)
		p.nodeData[node] = kvOffset
		m := p.nodeData[node+nVal]
		p.nodeData[node+nVal] = len(value)
		p.kvSize += len(value) - m
		return nil
	}

	h := p.randHeight()
	if h > p.maxHeight {
		for i := p.maxHeight; i < h; i++ {
			p.prevNode[i] = 0
		}
		p.maxHeight = h
	}

	kvOffset := len(p.kvData)
	p.kvData = append(p.kvData, key...)
	p.kvData = append(p.kvData, value...)
	// Node
	node := len(p.nodeData)
	p.nodeData = append(p.nodeData, kvOffset, len(key), len(value), h)
	for i, n := range p.prevNode[:h] {
		m := n + 4 + i
		p.nodeData = append(p.nodeData, p.nodeData[m])
		p.nodeData[m] = node
	}

	p.kvSize += len(key) + len(value)
	p.n++
	return nil
}

// Delete deletes the value for the given key. It returns ErrNotFound if
// the DB does not contain the key.
//
// It is safe to modify the contents of the arguments after Delete returns.
func (p *DB) Delete(key []byte) error {
	p.mu.Lock()
	defer p.mu.Unlock()

	node, exact := p.findGE(key, true)
	if !exact {
		return ErrNotFound
	}

	h := p.nodeData[node+nHeight]
	for i, n := range p.prevNode[:h] {
		m := n + 4 + i
		p.nodeData[m] = p.nodeData[p.nodeData[m]+nNext+i]
	}

	p.kvSize -= p.nodeData[node+nKey] + p.nodeData[node+nVal]
	p.n--
	return nil
}

// Contains returns true if the given key are in the DB.
//
// It is safe to modify the contents of the arguments after Contains returns.
func (p *DB) Contains(key []byte) bool {
	p.mu.RLock()
	_, exact := p.findGE(key, false)
	p.mu.RUnlock()
	return exact
}

// Get gets the value for the given key. It returns error.ErrNotFound if the
// DB does not contain the key.
//
// The caller should not modify the contents of the returned slice, but
// it is safe to modify the contents of the argument after Get returns.
func (p *DB) Get(key []byte) (value []byte, err error) {
	p.mu.RLock()
	if node, exact := p.findGE(key, false); exact {
		o := p.nodeData[node] + p.nodeData[node+nKey]
		value = p.kvData[o : o+p.nodeData[node+nVal]]
	} else {
		err = ErrNotFound
	}
	p.mu.RUnlock()
	return
}

// Find finds key/value pair whose key is greater than or equal to the
// given key. It returns ErrNotFound if the table doesn't contain
// such pair.
//
// The caller should not modify the contents of the returned slice, but
// it is safe to modify the contents of the argument after Find returns.
func (p *DB) Find(key []byte) (rkey, value []byte, err error) {
	p.mu.RLock()
	if node, _ := p.findGE(key, false); node != 0 {
		n := p.nodeData[node]
		m := n + p.nodeData[node+nKey]
		rkey = p.kvData[n:m]
		value = p.kvData[m : m+p.nodeData[node+nVal]]
	} else {
		err = ErrNotFound
	}
	p.mu.RUnlock()
	return
}

// NewIterator returns an iterator of the DB.
// The returned iterator is not goroutine-safe, but it is safe to use
// multiple iterators concurrently, with each in a dedicated goroutine.
// It is also safe to use an iterator concurrently with modifying its
// underlying DB. However, the resultant key/value pairs are not guaranteed
// to be a consistent snapshot of the DB at a particular point in time.
//
// Slice allows slicing the iterator to only contains keys in the given
// range. A nil Range.Start is treated as a key before all keys in the
// DB. And a nil Range.Limit is treated as a key after all keys in
// the DB.
//
// The iterator must be released after use, by calling Release method.
//
// Also read Iterator documentation of the leveldb/iterator package.
func (p *DB) NewIterator(slice *util.Range) iterator.Iterator {
	return &dbIter{p: p, slice: slice}
}

// Capacity returns keys/values buffer capacity.
func (p *DB) Capacity() int {
	p.mu.RLock()
	defer p.mu.RUnlock()
	return cap(p.kvData)
}

// Size returns sum of keys and values length. Note that deleted
// key/value will not be accouted for, but it will still consume
// the buffer, since the buffer is append only.
func (p *DB) Size() int {
	p.mu.RLock()
	defer p.mu.RUnlock()
	return p.kvSize
}

// Free returns keys/values free buffer before need to grow.
func (p *DB) Free() int {
	p.mu.RLock()
	defer p.mu.RUnlock()
	return cap(p.kvData) - len(p.kvData)
}

// Len returns the number of entries in the DB.
func (p *DB) Len() int {
	p.mu.RLock()
	defer p.mu.RUnlock()
	return p.n
}

// Reset resets the DB to initial empty state. Allows reuse the buffer.
func (p *DB) Reset() {
	p.rnd = rand.New(rand.NewSource(0xdeadbeef))
	p.maxHeight = 1
	p.n = 0
	p.kvSize = 0
	p.kvData = p.kvData[:0]
	p.nodeData = p.nodeData[:4+tMaxHeight]
	p.nodeData[nKV] = 0
	p.nodeData[nKey] = 0
	p.nodeData[nVal] = 0
	p.nodeData[nHeight] = tMaxHeight
	for n := 0; n < tMaxHeight; n++ {
		p.nodeData[4+n] = 0
		p.prevNode[n] = 0
	}
}

// New creates a new initalized in-memory key/value DB. The capacity
// is the initial key/value buffer capacity. The capacity is advisory,
// not enforced.
func New(cmp comparer.BasicComparer, capacity int) *DB {
	p := &DB{
		cmp:       cmp,
		rnd:       rand.New(rand.NewSource(0xdeadbeef)),
		maxHeight: 1,
		kvData:    make([]byte, 0, capacity),
		nodeData:  make([]int, 4+tMaxHeight),
	}
	p.nodeData[nHeight] = tMaxHeight
	return p
}
<commit_msg>memdb: Set variables to nil when iterator get released<commit_after>// Copyright (c) 2012, Suryandaru Triandana <syndtr@gmail.com>
// All rights reserved.
//
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

// Package memdb provides in-memory key/value database implementation.
package memdb

import (
	"math/rand"
	"sync"

	"github.com/conformal/goleveldb/leveldb/comparer"
	"github.com/conformal/goleveldb/leveldb/iterator"
	"github.com/conformal/goleveldb/leveldb/util"
)

var (
	ErrNotFound = util.ErrNotFound
)

const tMaxHeight = 12

type dbIter struct {
	util.BasicReleaser
	p          *DB
	slice      *util.Range
	node       int
	forward    bool
	key, value []byte
}

func (i *dbIter) fill(checkStart, checkLimit bool) bool {
	if i.node != 0 {
		n := i.p.nodeData[i.node]
		m := n + i.p.nodeData[i.node+nKey]
		i.key = i.p.kvData[n:m]
		if i.slice != nil {
			switch {
			case checkLimit && i.slice.Limit != nil && i.p.cmp.Compare(i.key, i.slice.Limit) >= 0:
				fallthrough
			case checkStart && i.slice.Start != nil && i.p.cmp.Compare(i.key, i.slice.Start) < 0:
				i.node = 0
				goto bail
			}
		}
		i.value = i.p.kvData[m : m+i.p.nodeData[i.node+nVal]]
		return true
	}
bail:
	i.key = nil
	i.value = nil
	return false
}

func (i *dbIter) Valid() bool {
	return i.node != 0
}

func (i *dbIter) First() bool {
	i.forward = true
	i.p.mu.RLock()
	defer i.p.mu.RUnlock()
	if i.slice != nil && i.slice.Start != nil {
		i.node, _ = i.p.findGE(i.slice.Start, false)
	} else {
		i.node = i.p.nodeData[nNext]
	}
	return i.fill(false, true)
}

func (i *dbIter) Last() bool {
	if i.p == nil {
		return false
	}
	i.forward = false
	i.p.mu.RLock()
	defer i.p.mu.RUnlock()
	if i.slice != nil && i.slice.Limit != nil {
		i.node = i.p.findLT(i.slice.Limit)
	} else {
		i.node = i.p.findLast()
	}
	return i.fill(true, false)
}

func (i *dbIter) Seek(key []byte) bool {
	if i.p == nil {
		return false
	}
	i.forward = true
	i.p.mu.RLock()
	defer i.p.mu.RUnlock()
	if i.slice != nil && i.slice.Start != nil && i.p.cmp.Compare(key, i.slice.Start) < 0 {
		key = i.slice.Start
	}
	i.node, _ = i.p.findGE(key, false)
	return i.fill(false, true)
}

func (i *dbIter) Next() bool {
	if i.p == nil {
		return false
	}
	if i.node == 0 {
		if !i.forward {
			return i.First()
		}
		return false
	}
	i.forward = true
	i.p.mu.RLock()
	defer i.p.mu.RUnlock()
	i.node = i.p.nodeData[i.node+nNext]
	return i.fill(false, true)
}

func (i *dbIter) Prev() bool {
	if i.p == nil {
		return false
	}
	if i.node == 0 {
		if i.forward {
			return i.Last()
		}
		return false
	}
	i.forward = false
	i.p.mu.RLock()
	defer i.p.mu.RUnlock()
	i.node = i.p.findLT(i.key)
	return i.fill(true, false)
}

func (i *dbIter) Key() []byte {
	return i.key
}

func (i *dbIter) Value() []byte {
	return i.value
}

func (i *dbIter) Error() error { return nil }

func (i *dbIter) Release() {
	if i.p != nil {
		i.p = nil
		i.node = 0
		i.key = nil
		i.value = nil
		i.BasicReleaser.Release()
	}
}

const (
	nKV = iota
	nKey
	nVal
	nHeight
	nNext
)

// DB is an in-memory key/value database.
type DB struct {
	cmp comparer.BasicComparer
	rnd *rand.Rand

	mu     sync.RWMutex
	kvData []byte
	// Node data:
	// [0]         : KV offset
	// [1]         : Key length
	// [2]         : Value length
	// [3]         : Height
	// [3..height] : Next nodes
	nodeData  []int
	prevNode  [tMaxHeight]int
	maxHeight int
	n         int
	kvSize    int
}

func (p *DB) randHeight() (h int) {
	const branching = 4
	h = 1
	for h < tMaxHeight && p.rnd.Int()%branching == 0 {
		h++
	}
	return
}

func (p *DB) findGE(key []byte, prev bool) (int, bool) {
	node := 0
	h := p.maxHeight - 1
	for {
		next := p.nodeData[node+nNext+h]
		cmp := 1
		if next != 0 {
			o := p.nodeData[next]
			cmp = p.cmp.Compare(p.kvData[o:o+p.nodeData[next+nKey]], key)
		}
		if cmp < 0 {
			// Keep searching in this list
			node = next
		} else {
			if prev {
				p.prevNode[h] = node
			} else if cmp == 0 {
				return next, true
			}
			if h == 0 {
				return next, cmp == 0
			}
			h--
		}
	}
}

func (p *DB) findLT(key []byte) int {
	node := 0
	h := p.maxHeight - 1
	for {
		next := p.nodeData[node+nNext+h]
		o := p.nodeData[next]
		if next == 0 || p.cmp.Compare(p.kvData[o:o+p.nodeData[next+nKey]], key) >= 0 {
			if h == 0 {
				break
			}
			h--
		} else {
			node = next
		}
	}
	return node
}

func (p *DB) findLast() int {
	node := 0
	h := p.maxHeight - 1
	for {
		next := p.nodeData[node+nNext+h]
		if next == 0 {
			if h == 0 {
				break
			}
			h--
		} else {
			node = next
		}
	}
	return node
}

// Put sets the value for the given key. It overwrites any previous value
// for that key; a DB is not a multi-map.
//
// It is safe to modify the contents of the arguments after Put returns.
func (p *DB) Put(key []byte, value []byte) error {
	p.mu.Lock()
	defer p.mu.Unlock()

	if node, exact := p.findGE(key, true); exact {
		kvOffset := len(p.kvData)
		p.kvData = append(p.kvData, key...)
		p.kvData = append(p.kvData, value...)
		p.nodeData[node] = kvOffset
		m := p.nodeData[node+nVal]
		p.nodeData[node+nVal] = len(value)
		p.kvSize += len(value) - m
		return nil
	}

	h := p.randHeight()
	if h > p.maxHeight {
		for i := p.maxHeight; i < h; i++ {
			p.prevNode[i] = 0
		}
		p.maxHeight = h
	}

	kvOffset := len(p.kvData)
	p.kvData = append(p.kvData, key...)
	p.kvData = append(p.kvData, value...)
	// Node
	node := len(p.nodeData)
	p.nodeData = append(p.nodeData, kvOffset, len(key), len(value), h)
	for i, n := range p.prevNode[:h] {
		m := n + 4 + i
		p.nodeData = append(p.nodeData, p.nodeData[m])
		p.nodeData[m] = node
	}

	p.kvSize += len(key) + len(value)
	p.n++
	return nil
}

// Delete deletes the value for the given key. It returns ErrNotFound if
// the DB does not contain the key.
//
// It is safe to modify the contents of the arguments after Delete returns.
func (p *DB) Delete(key []byte) error {
	p.mu.Lock()
	defer p.mu.Unlock()

	node, exact := p.findGE(key, true)
	if !exact {
		return ErrNotFound
	}

	h := p.nodeData[node+nHeight]
	for i, n := range p.prevNode[:h] {
		m := n + 4 + i
		p.nodeData[m] = p.nodeData[p.nodeData[m]+nNext+i]
	}

	p.kvSize -= p.nodeData[node+nKey] + p.nodeData[node+nVal]
	p.n--
	return nil
}

// Contains returns true if the given key are in the DB.
//
// It is safe to modify the contents of the arguments after Contains returns.
func (p *DB) Contains(key []byte) bool {
	p.mu.RLock()
	_, exact := p.findGE(key, false)
	p.mu.RUnlock()
	return exact
}

// Get gets the value for the given key. It returns error.ErrNotFound if the
// DB does not contain the key.
//
// The caller should not modify the contents of the returned slice, but
// it is safe to modify the contents of the argument after Get returns.
func (p *DB) Get(key []byte) (value []byte, err error) {
	p.mu.RLock()
	if node, exact := p.findGE(key, false); exact {
		o := p.nodeData[node] + p.nodeData[node+nKey]
		value = p.kvData[o : o+p.nodeData[node+nVal]]
	} else {
		err = ErrNotFound
	}
	p.mu.RUnlock()
	return
}

// Find finds key/value pair whose key is greater than or equal to the
// given key. It returns ErrNotFound if the table doesn't contain
// such pair.
//
// The caller should not modify the contents of the returned slice, but
// it is safe to modify the contents of the argument after Find returns.
func (p *DB) Find(key []byte) (rkey, value []byte, err error) {
	p.mu.RLock()
	if node, _ := p.findGE(key, false); node != 0 {
		n := p.nodeData[node]
		m := n + p.nodeData[node+nKey]
		rkey = p.kvData[n:m]
		value = p.kvData[m : m+p.nodeData[node+nVal]]
	} else {
		err = ErrNotFound
	}
	p.mu.RUnlock()
	return
}

// NewIterator returns an iterator of the DB.
// The returned iterator is not goroutine-safe, but it is safe to use
// multiple iterators concurrently, with each in a dedicated goroutine.
// It is also safe to use an iterator concurrently with modifying its
// underlying DB. However, the resultant key/value pairs are not guaranteed
// to be a consistent snapshot of the DB at a particular point in time.
//
// Slice allows slicing the iterator to only contains keys in the given
// range. A nil Range.Start is treated as a key before all keys in the
// DB. And a nil Range.Limit is treated as a key after all keys in
// the DB.
//
// The iterator must be released after use, by calling Release method.
//
// Also read Iterator documentation of the leveldb/iterator package.
func (p *DB) NewIterator(slice *util.Range) iterator.Iterator {
	return &dbIter{p: p, slice: slice}
}

// Capacity returns keys/values buffer capacity.
func (p *DB) Capacity() int {
	p.mu.RLock()
	defer p.mu.RUnlock()
	return cap(p.kvData)
}

// Size returns sum of keys and values length. Note that deleted
// key/value will not be accouted for, but it will still consume
// the buffer, since the buffer is append only.
func (p *DB) Size() int {
	p.mu.RLock()
	defer p.mu.RUnlock()
	return p.kvSize
}

// Free returns keys/values free buffer before need to grow.
func (p *DB) Free() int {
	p.mu.RLock()
	defer p.mu.RUnlock()
	return cap(p.kvData) - len(p.kvData)
}

// Len returns the number of entries in the DB.
func (p *DB) Len() int {
	p.mu.RLock()
	defer p.mu.RUnlock()
	return p.n
}

// Reset resets the DB to initial empty state. Allows reuse the buffer.
func (p *DB) Reset() {
	p.rnd = rand.New(rand.NewSource(0xdeadbeef))
	p.maxHeight = 1
	p.n = 0
	p.kvSize = 0
	p.kvData = p.kvData[:0]
	p.nodeData = p.nodeData[:4+tMaxHeight]
	p.nodeData[nKV] = 0
	p.nodeData[nKey] = 0
	p.nodeData[nVal] = 0
	p.nodeData[nHeight] = tMaxHeight
	for n := 0; n < tMaxHeight; n++ {
		p.nodeData[4+n] = 0
		p.prevNode[n] = 0
	}
}

// New creates a new initalized in-memory key/value DB. The capacity
// is the initial key/value buffer capacity. The capacity is advisory,
// not enforced.
func New(cmp comparer.BasicComparer, capacity int) *DB {
	p := &DB{
		cmp:       cmp,
		rnd:       rand.New(rand.NewSource(0xdeadbeef)),
		maxHeight: 1,
		kvData:    make([]byte, 0, capacity),
		nodeData:  make([]int, 4+tMaxHeight),
	}
	p.nodeData[nHeight] = tMaxHeight
	return p
}
<|endoftext|>