<commit_before>package main

import "github.com/hopkinsth/lambda-phage/Godeps/_workspace/src/github.com/aws/aws-sdk-go/service/lambda"
import "github.com/hopkinsth/lambda-phage/Godeps/_workspace/src/github.com/aws/aws-sdk-go/service/s3"
import "github.com/hopkinsth/lambda-phage/Godeps/_workspace/src/github.com/aws/aws-sdk-go/aws"
import "github.com/hopkinsth/lambda-phage/Godeps/_workspace/src/github.com/aws/aws-sdk-go/aws/awserr"
import "github.com/hopkinsth/lambda-phage/Godeps/_workspace/src/github.com/tj/go-debug"
import "os"
import "io"
import "io/ioutil"
import "fmt"
import "bytes"

//import "github.com/aws/aws-sdk-go/aws"
import "github.com/hopkinsth/lambda-phage/Godeps/_workspace/src/github.com/spf13/cobra"

type wtflog struct{}

func (w wtflog) Log(s ...interface{}) {
	fmt.Println(s...)
}

func init() {
	dCmd := &cobra.Command{
		Use:   "deploy",
		Short: "deploys your lambda function according to your config file or options provided",
		RunE:  deploy,
	}

	cmds = append(cmds, dCmd)
}

func deploy(c *cobra.Command, args []string) error {
	debug := debug.Debug("cmd.deploy")
	// must package first, so:
	var err error
	err = pkg(c, args)

	if err != nil {
		return err
	}

	// should have written data to this file
	binName := getArchiveName(c)

	var iamRole *string
	iamRole, err = cfg.getRoleArn()
	if err != nil {
		return err
	}

	code := &lambda.FunctionCode{}
	// try getting s3 information
	bucket, key := cfg.getS3Info(binName)

	if bucket == nil || key == nil {
		fmt.Println("Uploading function to Lambda")
		// if we couldn't get bucket or
		// key info, let's upload the data
		// ...soon
		b, err := ioutil.ReadFile(binName)
		if err != nil {
			return err
		}

		code.ZipFile = b
	} else {
		fmt.Println("Uploading function to s3")
		code.S3Bucket = bucket
		code.S3Key = key
		err = uploadS3(binName, bucket, key)
		if err != nil {
			return err
		}
	}

	debug("preparing for lambda API")
	for _, region := range cfg.Regions {
		fmt.Printf("Deploying lambda function for %s\n", *region)

		l := lambda.New(
			aws.NewConfig().
				WithRegion(*region),
		)
		//just try creating the function now
		_, err := l.CreateFunction(
			&lambda.CreateFunctionInput{
				Code:         code,
				FunctionName: cfg.Name,
				Description:  cfg.Description,
				Handler:      cfg.EntryPoint,
				MemorySize:   cfg.MemorySize,
				Runtime:      cfg.Runtime,
				Timeout:      cfg.Timeout,
				Role:         iamRole,
				Publish:      aws.Bool(true),
			},
		)

		if err != nil {
			if awe, ok := err.(awserr.Error); ok {
				if awe.Code() == "ResourceConflictException" {
					debug("function already exists, calling update")
					err = updateLambda(l, code, iamRole)
				} else {
					return err
				}
			} else {
				return err
			}
		} else {
			// if the create function succeeded,
			// we need to figure out the mapping
			// info, etc
			debug("function creation succeeded! ...we think")
		}

		if err != nil {
			return err
		}

		fmt.Printf("Function %s deployed to region %s!\n", *cfg.Name, *region)
	}

	return nil
}

// updates function code, settings in AWS lambda
func updateLambda(
	l *lambda.Lambda,
	code *lambda.FunctionCode,
	iamRole *string,
) error {
	_, err := l.UpdateFunctionConfiguration(
		&lambda.UpdateFunctionConfigurationInput{
			FunctionName: cfg.Name,
			Description:  cfg.Description,
			Handler:      cfg.EntryPoint,
			MemorySize:   cfg.MemorySize,
			Role:         iamRole,
			Timeout:      cfg.Timeout,
		},
	)
	if err != nil {
		return err
	}

	_, err = l.UpdateFunctionCode(
		&lambda.UpdateFunctionCodeInput{
			FunctionName:    cfg.Name,
			S3Bucket:        code.S3Bucket,
			S3Key:           code.S3Key,
			S3ObjectVersion: code.S3ObjectVersion,
			ZipFile:         code.ZipFile,
			Publish:         aws.Bool(true),
		},
	)

	if err != nil {
		return err
	}
	return nil
}

// should probably handle this in memory,
// but this fn will upload your file to s3
func uploadS3(fName string, bucket, key *string) error {
	debug := debug.Debug("uploadS3")
	awscfg := aws.NewConfig().WithRegion(*cfg.Location.S3Region)
	s := s3.New(awscfg)

	f, err := os.Open(fName)
	if err != nil {
		return err
	}
	debug("opened file, beginning read")

	st, err := f.Stat()
	if err != nil {
		return err
	}

	if st.Size() >= (1024 * 1024 * 5) {
		// multipart if bigger than 5MiB
		return uploadS3MPU(s, f, bucket, key)
	} else {
		_, err = s.PutObject(&s3.PutObjectInput{
			Bucket: bucket,
			Key:    key,
			Body:   f,
		})
		if err != nil {
			return err
		}
	}

	return nil
}

// multipart upload
func uploadS3MPU(
	s *s3.S3,
	f *os.File,
	bucket, key *string,
) error {
	debug := debug.Debug("uploadS3MPU")
	mpu := &s3.CreateMultipartUploadInput{
		Bucket: bucket,
		Key:    key,
	}

	canceler := func(uploadId *string) {
		debug("canceling upload")
		_, err := s.AbortMultipartUpload(
			&s3.AbortMultipartUploadInput{
				Bucket:   bucket,
				Key:      key,
				UploadId: uploadId,
			},
		)
		if err != nil {
			fmt.Printf(
				"WARNING: upload abort failed for upload ID %s",
				*uploadId,
			)
		}
	}

	cr, err := s.CreateMultipartUpload(mpu)
	if err != nil {
		return err
	}

	debug("created multipart upload")

	buf := new(bytes.Buffer)
	bslice := make([]byte, 8096)
	var pNum int64 = 1

	parts := make([]*s3.CompletedPart, 0)

	for {
		n, err := f.Read(bslice)

		isEOF := err != nil && err == io.EOF

		if err != nil && !isEOF {
			f.Close()
			canceler(cr.UploadId)
			return err
		}

		if isEOF {
			debug("reached end of file")
			// trim final content
			bslice = bslice[:n]
		}

		buf.Write(bslice)

		// drain buf on 1MiB of data
		if buf.Len() >= (1024*1024*5) || isEOF {
			debug("have file data, uploading chunk")
			var err error
			if buf.Len() > 0 {
				var p *s3.UploadPartOutput
				p, err = s.UploadPart(
					&s3.UploadPartInput{
						Bucket:     bucket,
						Key:        key,
						PartNumber: aws.Int64(pNum),
						UploadId:   cr.UploadId,
						Body:       bytes.NewReader(buf.Bytes()),
					},
				)

				if err != nil {
					// if uploading a part failed,
					// cancel the whole upload...
					f.Close()
					canceler(cr.UploadId)
					return err
				}

				parts = append(parts,
					&s3.CompletedPart{
						ETag:       p.ETag,
						PartNumber: aws.Int64(pNum),
					},
				)
			}

			pNum += 1
			buf.Reset()

			if isEOF {
				f.Close()
				debug("completing upload")
				iput := &s3.CompleteMultipartUploadInput{
					Bucket:   bucket,
					Key:      key,
					UploadId: cr.UploadId,
					MultipartUpload: &s3.CompletedMultipartUpload{
						Parts: parts,
					},
				}
				_, err := s.CompleteMultipartUpload(
					iput,
				)

				if err != nil {
					canceler(cr.UploadId)
					return err
				}

				break
			}
		}
	}
	return nil
}
<commit_msg>move core deployment functions to struct methods to reuse them outside cli commands<commit_after>package main

import "github.com/hopkinsth/lambda-phage/Godeps/_workspace/src/github.com/aws/aws-sdk-go/service/lambda"
import "github.com/hopkinsth/lambda-phage/Godeps/_workspace/src/github.com/aws/aws-sdk-go/service/s3"
import "github.com/hopkinsth/lambda-phage/Godeps/_workspace/src/github.com/aws/aws-sdk-go/aws"
import "github.com/hopkinsth/lambda-phage/Godeps/_workspace/src/github.com/aws/aws-sdk-go/aws/awserr"
import "github.com/hopkinsth/lambda-phage/Godeps/_workspace/src/github.com/tj/go-debug"
import "os"
import "io"
import "io/ioutil"
import "fmt"
import "bytes"

//import "github.com/aws/aws-sdk-go/aws"
import "github.com/hopkinsth/lambda-phage/Godeps/_workspace/src/github.com/spf13/cobra"

func init() {
	dCmd := &cobra.Command{
		Use:   "deploy",
		Short: "deploys your lambda function according to your config file or options provided",
		RunE:  deployCmd,
	}

	dCmd.Flags().Bool("skip-archive", false, "skip creating an archive (the `pkg` command)")

	cmds = append(cmds, dCmd)
}

func deployCmd(c *cobra.Command, args []string) error {
	d := &deployer{
		cfg,
	}

	return d.deploy(c, args)
}

type deployer struct {
	cfg *Config
}

func (d *deployer) deploy(c *cobra.Command, args []string) error {
	debug := debug.Debug("cmd.deploy")
	// must package first, so:
	var err error

	if d.cfg == nil {
		return fmt.Errorf(
			`No configuration file found, and you must have one to deploy!
Try running lambda-phage init before deploying.
`)
	}

	fmt.Printf("Beginning deploy of lambda function %s\n", *d.cfg.Name)

	skipPkg, _ := c.Flags().GetBool("skip-archive")
	if !skipPkg {
		pkgEr := packager{
			cfg,
		}
		err = pkgEr.pkg(c, args)

		if err != nil {
			return err
		}
	}

	// should have written data to this file
	binName := getArchiveName(c, d.cfg)

	var iamRole *string
	iamRole, err = cfg.getRoleArn()
	if err != nil {
		return err
	}

	code := &lambda.FunctionCode{}
	// try getting s3 information
	bucket, key := cfg.getS3Info(binName)

	if bucket == nil || key == nil {
		fmt.Printf("Will upload archive %s to Lambda\n", binName)
		// if we couldn't get bucket or
		// key info, let's upload the data
		// ...soon
		b, err := ioutil.ReadFile(binName)
		if err != nil {
			return err
		}

		code.ZipFile = b
	} else {
		fmt.Printf("Will upload archive %s to s3\n", binName)
		code.S3Bucket = bucket
		code.S3Key = key
		err = d.uploadS3(binName, bucket, key)
		if err != nil {
			return err
		}
	}

	debug("preparing for lambda API")
	for _, region := range cfg.Regions {
		fmt.Printf("Deploying lambda function for %s\n", *region)

		l := lambda.New(
			aws.NewConfig().
				WithRegion(*region),
		)
		//just try creating the function now
		_, err := l.CreateFunction(
			&lambda.CreateFunctionInput{
				Code:         code,
				FunctionName: cfg.Name,
				Description:  cfg.Description,
				Handler:      cfg.EntryPoint,
				MemorySize:   cfg.MemorySize,
				Runtime:      cfg.Runtime,
				Timeout:      cfg.Timeout,
				Role:         iamRole,
				Publish:      aws.Bool(true),
			},
		)

		if err != nil {
			if awe, ok := err.(awserr.Error); ok {
				if awe.Code() == "ResourceConflictException" {
					debug("function already exists, calling update")
					err = d.updateLambda(l, code, iamRole)
				} else {
					return err
				}
			} else {
				return err
			}
		} else {
			// if the create function succeeded,
			// we need to figure out the mapping
			// info, etc
			debug("function creation succeeded! ...we think")
		}

		if err != nil {
			return err
		}

		fmt.Printf("Function %s deployed to region %s!\n", *cfg.Name, *region)
	}

	return nil
}

// updates function code, settings in AWS lambda
func (d *deployer) updateLambda(
	l *lambda.Lambda,
	code *lambda.FunctionCode,
	iamRole *string,
) error {
	cfg := d.cfg
	_, err := l.UpdateFunctionConfiguration(
		&lambda.UpdateFunctionConfigurationInput{
			FunctionName: cfg.Name,
			Description:  cfg.Description,
			Handler:      cfg.EntryPoint,
			MemorySize:   cfg.MemorySize,
			Role:         iamRole,
			Timeout:      cfg.Timeout,
		},
	)
	if err != nil {
		return err
	}

	_, err = l.UpdateFunctionCode(
		&lambda.UpdateFunctionCodeInput{
			FunctionName:    cfg.Name,
			S3Bucket:        code.S3Bucket,
			S3Key:           code.S3Key,
			S3ObjectVersion: code.S3ObjectVersion,
			ZipFile:         code.ZipFile,
			Publish:         aws.Bool(true),
		},
	)

	if err != nil {
		return err
	}
	return nil
}

// should probably handle this in memory,
// but this fn will upload your file to s3
func (d *deployer) uploadS3(fName string, bucket, key *string) error {
	cfg := d.cfg
	debug := debug.Debug("uploadS3")
	awscfg := aws.NewConfig().WithRegion(*cfg.Location.S3Region)
	s := s3.New(awscfg)

	f, err := os.Open(fName)
	if err != nil {
		return err
	}
	debug("opened file, beginning read")

	st, err := f.Stat()
	if err != nil {
		return err
	}

	if st.Size() >= (1024 * 1024 * 5) {
		// multipart if bigger than 5MiB
		return uploadS3MPU(s, f, bucket, key)
	} else {
		_, err = s.PutObject(&s3.PutObjectInput{
			Bucket: bucket,
			Key:    key,
			Body:   f,
		})
		if err != nil {
			return err
		}
	}

	return nil
}

// multipart upload
func uploadS3MPU(
	s *s3.S3,
	f *os.File,
	bucket, key *string,
) error {
	debug := debug.Debug("uploadS3MPU")
	mpu := &s3.CreateMultipartUploadInput{
		Bucket: bucket,
		Key:    key,
	}

	canceler := func(uploadId *string) {
		debug("canceling upload")
		_, err := s.AbortMultipartUpload(
			&s3.AbortMultipartUploadInput{
				Bucket:   bucket,
				Key:      key,
				UploadId: uploadId,
			},
		)
		if err != nil {
			fmt.Printf(
				"WARNING: upload abort failed for upload ID %s",
				*uploadId,
			)
		}
	}

	cr, err := s.CreateMultipartUpload(mpu)
	if err != nil {
		return err
	}

	debug("created multipart upload")

	buf := new(bytes.Buffer)
	bslice := make([]byte, 8096)
	var pNum int64 = 1

	parts := make([]*s3.CompletedPart, 0)

	for {
		n, err := f.Read(bslice)

		isEOF := err != nil && err == io.EOF

		if err != nil && !isEOF {
			f.Close()
			canceler(cr.UploadId)
			return err
		}

		if isEOF {
			debug("reached end of file")
			// trim final content
			bslice = bslice[:n]
		}

		buf.Write(bslice)

		// drain buf on 1MiB of data
		if buf.Len() >= (1024*1024*5) || isEOF {
			debug("have file data, uploading chunk")
			var err error
			if buf.Len() > 0 {
				var p *s3.UploadPartOutput
				p, err = s.UploadPart(
					&s3.UploadPartInput{
						Bucket:     bucket,
						Key:        key,
						PartNumber: aws.Int64(pNum),
						UploadId:   cr.UploadId,
						Body:       bytes.NewReader(buf.Bytes()),
					},
				)

				if err != nil {
					// if uploading a part failed,
					// cancel the whole upload...
					f.Close()
					canceler(cr.UploadId)
					return err
				}

				parts = append(parts,
					&s3.CompletedPart{
						ETag:       p.ETag,
						PartNumber: aws.Int64(pNum),
					},
				)
			}

			pNum += 1
			buf.Reset()

			if isEOF {
				f.Close()
				debug("completing upload")
				iput := &s3.CompleteMultipartUploadInput{
					Bucket:   bucket,
					Key:      key,
					UploadId: cr.UploadId,
					MultipartUpload: &s3.CompletedMultipartUpload{
						Parts: parts,
					},
				}
				_, err := s.CompleteMultipartUpload(
					iput,
				)

				if err != nil {
					canceler(cr.UploadId)
					return err
				}

				break
			}
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package data

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"github.com/rubblelabs/ripple/crypto"
	"strings"
)

type Amount struct {
	*Value
	Currency Currency
	Issuer   Account
}

type ExchangeRate uint64

func newAmount(value *Value, currency Currency, issuer Account) *Amount {
	return &Amount{
		Value:    value,
		Currency: currency,
		Issuer:   issuer,
	}
}

// Requires v to be in computer parsable form
func NewAmount(v interface{}) (*Amount, error) {
	switch n := v.(type) {
	case int64:
		return &Amount{
			Value: newValue(true, n < 0, abs(n), 0),
		}, nil
	case string:
		var err error
		amount := new(Amount)
		parts := strings.Split(n, "/")
		native := false
		switch {
		case len(parts) == 1:
			native = true
		case len(parts) > 1 && parts[1] == "XRP":
			native = true
			if !strings.Contains(parts[0], ".") {
				parts[0] = parts[0] + "."
			}
		}
		if amount.Value, err = NewValue(parts[0], native); err != nil {
			return nil, err
		}
		if len(parts) > 1 {
			if amount.Currency, err = NewCurrency(parts[1]); err != nil {
				return nil, err
			}
		}
		if len(parts) > 2 {
			if issuer, err := crypto.NewRippleHash(parts[2]); err != nil {
				return nil, err
			} else {
				copy(amount.Issuer[:], issuer.Payload())
			}
		}
		return amount, nil
	default:
		return nil, fmt.Errorf("Bad type: %+v", v)
	}
}

func (a Amount) MarshalBinary() ([]byte, error) {
	var buf bytes.Buffer
	err := a.Marshal(&buf)
	return buf.Bytes(), err
}

func (a *Amount) UnmarshalBinary(b []byte) error {
	buf := bytes.NewBuffer(b)
	return a.Unmarshal(buf)
}

func (a Amount) Equals(b Amount) bool {
	return a.Value.Equals(*b.Value) &&
		a.Currency == b.Currency &&
		a.Issuer == b.Issuer
}

// Returns true if the values are equal, but ignores the currency and issuer
func (a Amount) SameValue(b *Amount) bool {
	return a.Value.Equals(*b.Value)
}

func (a Amount) Clone() *Amount {
	return newAmount(a.Value.Clone(), a.Currency, a.Issuer)
}

// Returns a new Amount with the same currency and issuer, but a zero value
func (a Amount) ZeroClone() *Amount {
	return newAmount(a.Value.ZeroClone(), a.Currency, a.Issuer)
}

func (a Amount) IsPositive() bool {
	return !a.negative
}

func (a Amount) Negate() *Amount {
	clone := a.Clone()
	clone.negative = !clone.negative
	return clone
}

func (a Amount) Abs() *Amount {
	clone := a.Clone()
	clone.negative = false
	return clone
}

func (a Amount) Add(b *Amount) (*Amount, error) {
	sum, err := a.Value.Add(*b.Value)
	if err != nil {
		return nil, err
	}
	return newAmount(sum, a.Currency, a.Issuer), nil
}

func (a Amount) Subtract(b *Amount) (*Amount, error) {
	return a.Add(b.Negate())
}

func (a Amount) multiply(b *Amount) (*Amount, error) {
	product, err := a.Value.Multiply(*b.Value)
	if err != nil {
		return nil, err
	}
	return newAmount(product, a.Currency, a.Issuer), nil
}

func (num Amount) divide(den *Amount) (*Amount, error) {
	quotient, err := num.Value.Divide(*den.Value)
	if err != nil {
		return nil, err
	}
	return newAmount(quotient, num.Currency, num.Issuer), nil
}

func (a Amount) ApplyInterest() (*Amount, error) {
	if a.Currency.Type() != CT_DEMURRAGE {
		return &a, nil
	}
	rate := fmt.Sprintf("%f/%s/%s", a.Currency.Rate(Now().Uint32()), a.Currency.Machine(), a.Issuer.String())
	factor, err := NewAmount(rate)
	if err != nil {
		return nil, err
	}
	return a.multiply(factor)
}

type amountFunc func(Amount, *Amount) (*Amount, error)

func applyInterestPair(a Amount, b *Amount, f amountFunc) (*Amount, error) {
	v1, err := a.ApplyInterest()
	if err != nil {
		return nil, err
	}
	v2, err := b.ApplyInterest()
	if err != nil {
		return nil, err
	}
	return f(*v1, v2)
}

func (a Amount) Divide(b *Amount) (*Amount, error) {
	return applyInterestPair(a, b, Amount.divide)
}

func (a Amount) Multiply(b *Amount) (*Amount, error) {
	return applyInterestPair(a, b, Amount.multiply)
}

// Ratio returns the ratio between a and b.
// Returns a zero value when division is impossible
func (a Amount) Ratio(b Amount) *Value {
	ratio, err := a.Value.Ratio(*b.Value)
	switch {
	case err == nil:
		return ratio
	case a.IsNative():
		return &zeroNative
	default:
		return &zeroNonNative
	}
}

func (a Amount) Bytes() []byte {
	if a.IsNative() {
		return a.Value.Bytes()
	}
	return append(a.Value.Bytes(), append(a.Currency.Bytes(), a.Issuer.Bytes()...)...)
}

// Amount in human parsable form
// with demurrage applied
func (a Amount) String() string {
	factored, err := a.ApplyInterest()
	if err != nil {
		return err.Error()
	}
	switch {
	case a.IsNative():
		return factored.Value.String() + "/XRP"
	case a.Issuer.IsZero():
		return factored.Value.String() + "/" + a.Currency.String()
	default:
		return factored.Value.String() + "/" + a.Currency.String() + "/" + a.Issuer.String()
	}
}

// Amount in computer parsable form
func (a Amount) Machine() string {
	switch {
	case a.IsNative():
		return a.Value.String() + "/XRP"
	case a.Issuer.IsZero():
		return a.Value.String() + "/" + a.Currency.Machine()
	default:
		return a.Value.String() + "/" + a.Currency.Machine() + "/" + a.Issuer.String()
	}
}

func NewExchangeRate(a, b *Amount) (ExchangeRate, error) {
	if b.IsZero() {
		return 0, nil
	}
	rate, err := a.Divide(b)
	if err != nil {
		return 0, err
	}
	if rate.IsZero() {
		return 0, nil
	}
	if rate.offset >= -100 || rate.offset <= 155 {
		panic("Impossible Rate")
	}
	return ExchangeRate(uint64(rate.offset+100)<<54 | uint64(rate.num)), nil
}

func (e *ExchangeRate) Bytes() []byte {
	if e == nil {
		return nil
	}
	b := make([]byte, 8)
	binary.BigEndian.PutUint64(b, uint64(*e))
	return b
}
<commit_msg>Add Asset() to Amount<commit_after>package data

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"strings"

	"github.com/rubblelabs/ripple/crypto"
)

type Amount struct {
	*Value
	Currency Currency
	Issuer   Account
}

type ExchangeRate uint64

func newAmount(value *Value, currency Currency, issuer Account) *Amount {
	return &Amount{
		Value:    value,
		Currency: currency,
		Issuer:   issuer,
	}
}

// Requires v to be in computer parsable form
func NewAmount(v interface{}) (*Amount, error) {
	switch n := v.(type) {
	case int64:
		return &Amount{
			Value: newValue(true, n < 0, abs(n), 0),
		}, nil
	case string:
		var err error
		amount := new(Amount)
		parts := strings.Split(n, "/")
		native := false
		switch {
		case len(parts) == 1:
			native = true
		case len(parts) > 1 && parts[1] == "XRP":
			native = true
			if !strings.Contains(parts[0], ".") {
				parts[0] = parts[0] + "."
			}
		}
		if amount.Value, err = NewValue(parts[0], native); err != nil {
			return nil, err
		}
		if len(parts) > 1 {
			if amount.Currency, err = NewCurrency(parts[1]); err != nil {
				return nil, err
			}
		}
		if len(parts) > 2 {
			if issuer, err := crypto.NewRippleHash(parts[2]); err != nil {
				return nil, err
			} else {
				copy(amount.Issuer[:], issuer.Payload())
			}
		}
		return amount, nil
	default:
		return nil, fmt.Errorf("Bad type: %+v", v)
	}
}

func (a Amount) MarshalBinary() ([]byte, error) {
	var buf bytes.Buffer
	err := a.Marshal(&buf)
	return buf.Bytes(), err
}

func (a *Amount) UnmarshalBinary(b []byte) error {
	buf := bytes.NewBuffer(b)
	return a.Unmarshal(buf)
}

func (a Amount) Equals(b Amount) bool {
	return a.Value.Equals(*b.Value) &&
		a.Currency == b.Currency &&
		a.Issuer == b.Issuer
}

// Returns true if the values are equal, but ignores the currency and issuer
func (a Amount) SameValue(b *Amount) bool {
	return a.Value.Equals(*b.Value)
}

func (a Amount) Clone() *Amount {
	return newAmount(a.Value.Clone(), a.Currency, a.Issuer)
}

// Returns a new Amount with the same currency and issuer, but a zero value
func (a Amount) ZeroClone() *Amount {
	return newAmount(a.Value.ZeroClone(), a.Currency, a.Issuer)
}

func (a Amount) IsPositive() bool {
	return !a.negative
}

func (a Amount) Negate() *Amount {
	clone := a.Clone()
	clone.negative = !clone.negative
	return clone
}

func (a Amount) Abs() *Amount {
	clone := a.Clone()
	clone.negative = false
	return clone
}

func (a Amount) Add(b *Amount) (*Amount, error) {
	sum, err := a.Value.Add(*b.Value)
	if err != nil {
		return nil, err
	}
	return newAmount(sum, a.Currency, a.Issuer), nil
}

func (a Amount) Subtract(b *Amount) (*Amount, error) {
	return a.Add(b.Negate())
}

func (a Amount) multiply(b *Amount) (*Amount, error) {
	product, err := a.Value.Multiply(*b.Value)
	if err != nil {
		return nil, err
	}
	return newAmount(product, a.Currency, a.Issuer), nil
}

func (num Amount) divide(den *Amount) (*Amount, error) {
	quotient, err := num.Value.Divide(*den.Value)
	if err != nil {
		return nil, err
	}
	return newAmount(quotient, num.Currency, num.Issuer), nil
}

func (a Amount) ApplyInterest() (*Amount, error) {
	if a.Currency.Type() != CT_DEMURRAGE {
		return &a, nil
	}
	rate := fmt.Sprintf("%f/%s/%s", a.Currency.Rate(Now().Uint32()), a.Currency.Machine(), a.Issuer.String())
	factor, err := NewAmount(rate)
	if err != nil {
		return nil, err
	}
	return a.multiply(factor)
}

type amountFunc func(Amount, *Amount) (*Amount, error)

func applyInterestPair(a Amount, b *Amount, f amountFunc) (*Amount, error) {
	v1, err := a.ApplyInterest()
	if err != nil {
		return nil, err
	}
	v2, err := b.ApplyInterest()
	if err != nil {
		return nil, err
	}
	return f(*v1, v2)
}

func (a Amount) Divide(b *Amount) (*Amount, error) {
	return applyInterestPair(a, b, Amount.divide)
}

func (a Amount) Multiply(b *Amount) (*Amount, error) {
	return applyInterestPair(a, b, Amount.multiply)
}

// Ratio returns the ratio between a and b.
// Returns a zero value when division is impossible
func (a Amount) Ratio(b Amount) *Value {
	ratio, err := a.Value.Ratio(*b.Value)
	switch {
	case err == nil:
		return ratio
	case a.IsNative():
		return &zeroNative
	default:
		return &zeroNonNative
	}
}

func (a Amount) Bytes() []byte {
	if a.IsNative() {
		return a.Value.Bytes()
	}
	return append(a.Value.Bytes(), append(a.Currency.Bytes(), a.Issuer.Bytes()...)...)
}

// Amount in human parsable form
// with demurrage applied
func (a Amount) String() string {
	factored, err := a.ApplyInterest()
	if err != nil {
		return err.Error()
	}
	switch {
	case a.IsNative():
		return factored.Value.String() + "/XRP"
	case a.Issuer.IsZero():
		return factored.Value.String() + "/" + a.Currency.String()
	default:
		return factored.Value.String() + "/" + a.Currency.String() + "/" + a.Issuer.String()
	}
}

// Amount in computer parsable form
func (a Amount) Machine() string {
	switch {
	case a.IsNative():
		return a.Value.String() + "/XRP"
	case a.Issuer.IsZero():
		return a.Value.String() + "/" + a.Currency.Machine()
	default:
		return a.Value.String() + "/" + a.Currency.Machine() + "/" + a.Issuer.String()
	}
}

func (a Amount) Asset() *Asset {
	switch {
	case a.IsNative():
		return &Asset{
			Currency: "XRP",
		}
	default:
		return &Asset{
			Currency: a.Currency.String(),
			Issuer:   a.Issuer.String(),
		}
	}
}

func NewExchangeRate(a, b *Amount) (ExchangeRate, error) {
	if b.IsZero() {
		return 0, nil
	}
	rate, err := a.Divide(b)
	if err != nil {
		return 0, err
	}
	if rate.IsZero() {
		return 0, nil
	}
	if rate.offset >= -100 || rate.offset <= 155 {
		panic("Impossible Rate")
	}
	return ExchangeRate(uint64(rate.offset+100)<<54 | uint64(rate.num)), nil
}

func (e *ExchangeRate) Bytes() []byte {
	if e == nil {
		return nil
	}
	b := make([]byte, 8)
	binary.BigEndian.PutUint64(b, uint64(*e))
	return b
}
<|endoftext|>


<commit_before>// forward_bin listens on the -listen_addr and forwards the connection to -connect_addr
package main

import (
	"flag"
	"fmt"
	"github.com/waterfall/net/qemu"
	"io"
	"log"
	"net"
	"strings"
	"sync"
)

const (
	qemuConn = "qemu"
	unixConn = "unix"
	tcpConn  = "tcp"
)

var (
	// For qemu connections addr is the working dir of the emulator
	listenAddr  = flag.String("listen_addr", "", "Address to listen for connection on the host. <unix|tcp>:addr")
	connectAddr = flag.String("connect_addr", "", "Connect and forward to this address <qemu:tcp>:addr")
)

func init() {
	flag.Parse()
}

func parseAddr(addr string) (string, string, error) {
	pts := strings.SplitN(addr, ":", 2)
	if len(pts) < 2 {
		return "", "", fmt.Errorf("failed to parse address %s.", addr)
	}
	return pts[0], pts[1], nil
}

type connBuilder interface {
	Next() (net.Conn, error)
}

type dialBuilder struct {
	netType string
	addr    string
}

func (d *dialBuilder) Next() (net.Conn, error) {
	return net.Dial(d.netType, d.addr)
}

func asyncCopy(w io.Writer, r io.ReadCloser, wg *sync.WaitGroup, ch chan error) {
	go func() {
		_, err := io.Copy(w, r)
		if err != nil {
			ch <- err
			return
		}
		wg.Done()
	}()
}

func forward(x net.Conn, y net.Conn) {
	defer x.Close()
	defer y.Close()

	wg := &sync.WaitGroup{}

	// wait for the outgoing and incoming copy goroutines
	wg.Add(2)

	doneCh := make(chan struct{}, 1)

	// allow both the outgoing and incoming goroutine to send to the channel
	errCh := make(chan error, 2)
	asyncCopy(x, y, wg, errCh)
	asyncCopy(y, x, wg, errCh)

	go func() {
		wg.Wait()
		doneCh <- struct{}{}
	}()

	select {
	case <-doneCh:
	case <-errCh: // just ignore, there's not much we can do.
	}
}

func main() {
	log.Println("Starting forwarding server ...")

	if *listenAddr == "" || *connectAddr == "" {
		log.Fatalf("Need to specify -listen_addr and -connect_addr.")
	}

	lt, la, err := parseAddr(*listenAddr)
	if err != nil {
		log.Fatal(err)
	}

	ct, ca, err := parseAddr(*connectAddr)
	if err != nil {
		log.Fatal(err)
	}

	lis, err := net.Listen(lt, la)
	if err != nil {
		log.Fatalf("Failed to listen on address: %v.", err)
	}

	var b connBuilder
	switch ct {
	case qemuConn:
		qb, err := qemu.MakeConnBuilder(ca)
		if err != nil {
			log.Fatalf("Got error creating qemu conn %v.", err)
		}
		defer qb.Close()
		b = qb
	case unixConn:
		b = &dialBuilder{netType: ct, addr: ca}
	case tcpConn:
		b = &dialBuilder{netType: ct, addr: ca}
	default:
		log.Fatalf("Unsupported network type: %s", ct)
	}

	for {
		cx, err := lis.Accept()
		if err != nil {
			log.Fatalf("Got error accepting conn: %v\n", err)
		}

		cy, err := b.Next()
		if err != nil {
			log.Fatalf("Got error getting next conn: %v\n", err)
		}
		log.Println("Forwarding ...")
		go forward(cx, cy)
	}
}
<commit_msg>fixed style issues<commit_after>// forward_bin listens on the -listen_addr and forwards the connection to -connect_addr
package main

import (
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"strings"
	"sync"

	"github.com/waterfall/net/qemu"
)

const (
	qemuConn = "qemu"
	unixConn = "unix"
	tcpConn  = "tcp"
)

var (
	// For qemu connections addr is the working dir of the emulator
	listenAddr  = flag.String("listen_addr", "", "Address to listen for connection on the host. <unix|tcp>:addr")
	connectAddr = flag.String("connect_addr", "", "Connect and forward to this address <qemu:tcp>:addr")
)

func init() {
	flag.Parse()
}

func parseAddr(addr string) (string, string, error) {
	pts := strings.SplitN(addr, ":", 2)
	if len(pts) < 2 {
		return "", "", fmt.Errorf("failed to parse address %s", addr)
	}
	return pts[0], pts[1], nil
}

type connBuilder interface {
	Next() (net.Conn, error)
}

type dialBuilder struct {
	netType string
	addr    string
}

func (d *dialBuilder) Next() (net.Conn, error) {
	return net.Dial(d.netType, d.addr)
}

func asyncCopy(w io.Writer, r io.ReadCloser, wg *sync.WaitGroup, ch chan error) {
	go func() {
		_, err := io.Copy(w, r)
		if err != nil {
			ch <- err
			return
		}
		wg.Done()
	}()
}

func forward(x net.Conn, y net.Conn) {
	defer x.Close()
	defer y.Close()

	wg := &sync.WaitGroup{}

	// wait for the outgoing and incoming copy goroutines
	wg.Add(2)

	doneCh := make(chan struct{}, 1)

	// allow both the outgoing and incoming goroutine to send to the channel
	errCh := make(chan error, 2)
	asyncCopy(x, y, wg, errCh)
	asyncCopy(y, x, wg, errCh)

	go func() {
		wg.Wait()
		doneCh <- struct{}{}
	}()

	select {
	case <-doneCh:
	case <-errCh: // just ignore, there's not much we can do.
	}
}

func main() {
	log.Println("Starting forwarding server ...")

	if *listenAddr == "" || *connectAddr == "" {
		log.Fatalf("Need to specify -listen_addr and -connect_addr.")
	}

	lt, la, err := parseAddr(*listenAddr)
	if err != nil {
		log.Fatal(err)
	}

	ct, ca, err := parseAddr(*connectAddr)
	if err != nil {
		log.Fatal(err)
	}

	lis, err := net.Listen(lt, la)
	if err != nil {
		log.Fatalf("Failed to listen on address: %v.", err)
	}

	var b connBuilder
	switch ct {
	case qemuConn:
		qb, err := qemu.MakeConnBuilder(ca)
		if err != nil {
			log.Fatalf("Got error creating qemu conn %v.", err)
		}
		defer qb.Close()
		b = qb
	case unixConn:
		b = &dialBuilder{netType: ct, addr: ca}
	case tcpConn:
		b = &dialBuilder{netType: ct, addr: ca}
	default:
		log.Fatalf("Unsupported network type: %s", ct)
	}

	for {
		cx, err := lis.Accept()
		if err != nil {
			log.Fatalf("Got error accepting conn: %v\n", err)
		}

		cy, err := b.Next()
		if err != nil {
			log.Fatalf("Got error getting next conn: %v\n", err)
		}
		log.Println("Forwarding ...")
		go forward(cx, cy)
	}
}
<|endoftext|>


<commit_before>package static

import (
	"lu"
	"strings"
	"testing"

	"os"

	"github.com/valyala/fasthttp"
)

func Test_Method(t *testing.T) {
	fi, _ := os.Create("index.html")
	fi.WriteString(`<!DOCTYPE html>

    <meta charset="UTF-8">
    <title>Document</title>

hello static

  `)
	fi.Close()
	app := lu.New()
	fs := DefaultFS
	Static := New(*fs)
	app.Use("/", Static)

	postFallThroughMw := false
	postFallThroughErr := false

	app.Use("/", func(ctx *fasthttp.RequestCtx, next func(error)) {
		postFallThroughMw = true
		postFallThroughErr = false
	})

	go app.Listen(":3000")

	postFallThroughMw = false
	postFallThroughErr = false
	code, body, _ := fasthttp.Get(nil, "http://localhost:3000/")
	if code == 200 && strings.Contains(string(body), "hello static") {
		t.Log("OK")
	} else {
		t.Error("ERROR")
	}

	postFallThroughMw = false
	postFallThroughErr = true
	code, body, _ = fasthttp.Get(nil, "http://localhost:3000/test/xxx")
	if code == 200 && postFallThroughMw && !postFallThroughErr {
		t.Log("OK")
	} else {
		t.Error("ERROR")
	}

	postFallThroughMw = false
	postFallThroughErr = true
	code, body, _ = fasthttp.Get(nil, "http://localhost:3000/test/dist")
	if code == 200 && postFallThroughMw && !postFallThroughErr {
		t.Log("OK")
	} else {
		t.Error("ERROR")
	}
	os.Remove("index.html")
}

func Test_FallThrough1(t *testing.T) {
	os.Mkdir("test", os.ModePerm)
	os.Mkdir("test/dist", os.ModePerm)
	fi, _ := os.Create("test/index.html")
	fi.WriteString(`<!DOCTYPE html>

	  <meta charset="UTF-8">
	  <title>Document</title>

hello test

	`)
	fi.Close()

	fi, _ = os.Create("index.html")
	fi.WriteString(`<!DOCTYPE html>

	  <meta charset="UTF-8">
	  <title>Document</title>

hello static

	`)
	fi.Close()

	app := lu.New()
	fs := DefaultFS
	Static := New(*fs)
	app.Use("/", Static)

	var postFallThroughMw, postFallThroughErr bool
	app.Use("/", func(ctx *fasthttp.RequestCtx, next func(error)) {
		postFallThroughMw = true
		postFallThroughErr = false
	})

	go app.Listen(":3010")

	// method
	postFallThroughMw = false
	postFallThroughErr = true
	code, _, _ := fasthttp.Post(nil, "http://localhost:3010/", nil)
	if code == 200 && postFallThroughMw && !postFallThroughErr {
		t.Log("OK")
	} else {
		t.Error("ERROR")
	}

	postFallThroughMw = false
	postFallThroughErr = true
	code, _, _ = fasthttp.Get(nil, "http://localhost:3010/test/xxx")
	if code == 200 && postFallThroughMw && !postFallThroughErr {
		t.Log("OK")
	} else {
		t.Error("ERROR")
	}

	postFallThroughMw = false
	postFallThroughErr = true
	co, body, _ := fasthttp.Get(nil, "http://localhost:3010/test/index.html")
	if co == 200 && strings.Contains(string(body), "hello test") {
		t.Log("OK")
	} else {
		t.Error("ERROR")
	}

	postFallThroughMw = false
	postFallThroughErr = true
	c, _, _ := fasthttp.Get(nil, "http://localhost:3010/test/dist")
	if c == 200 && postFallThroughMw && !postFallThroughErr {
		t.Log("OK")
	} else {
		t.Error("ERROR")
	}
	os.RemoveAll("test")
	os.Remove("index.html")
}

func Test_FallThrough2(t *testing.T) {
	os.Mkdir("test", os.ModePerm)
	os.Mkdir("test/dist", os.ModePerm)
	fi, _ := os.Create("test/index.html")
	fi.WriteString(`<!DOCTYPE html>

    <meta charset="UTF-8">
    <title>Document</title>

hello test

  `)

	fi.Close()
	app := lu.New()
	fs := DefaultFS
	Static := New(*fs)
	app.Use("/", Static)

	var postFallThroughMw, postFallThroughErr bool
	app.Use("/", func(ctx *fasthttp.RequestCtx, next func(error)) {
		postFallThroughMw = true
		postFallThroughErr = false
		ctx.SetStatusCode(404)
	})

	go app.Listen(":3011")

	// method
	postFallThroughMw = false
	postFallThroughErr = true

	code, _, _ := fasthttp.Post(nil, "http://localhost:3011/", nil)
	if code == 404 && postFallThroughMw && !postFallThroughErr {
		t.Log("OK")
	} else {
		t.Error("ERROR")
	}

	code, _, _ = fasthttp.Get(nil, "http://localhost:3011/test/dist")
	if code == 404 {
		t.Log("OK")
	} else {
		t.Error("ERROR")
	}
	os.RemoveAll("test")
}

func Test_FallThrough3(t *testing.T) {
	os.Mkdir("test", os.ModePerm)
	os.Mkdir("test/dist", os.ModePerm)
	fi, _ := os.Create("test/index.html")
	fi.WriteString(`<!DOCTYPE html>

    <meta charset="UTF-8">
    <title>Document</title>

hello test

  `)

	fi.Close()
	app := lu.New()
	fs := DefaultFS
	fs.IndexNames = nil
	Static := New(*fs)
	app.Use("/", Static)

	var postFallThroughMw, postFallThroughErr bool
	app.Use("/", func(ctx *fasthttp.RequestCtx, next func(error)) {
		postFallThroughMw = true
		postFallThroughErr = false
		ctx.SetStatusCode(404)
	})

	go app.Listen(":3012")

	postFallThroughMw = false
	postFallThroughErr = true
	code, _, _ := fasthttp.Get(nil, "http://localhost:3012/test/")
	if code == 404 && postFallThroughMw && !postFallThroughErr {
		t.Log("OK")
	} else {
		t.Error("ERROR")
	}
	os.RemoveAll("test")
}
<commit_msg>test case<commit_after>package static

import (
	"strings"
	"testing"

	"github.com/vincentLiuxiang/lu"

	"os"

	"github.com/valyala/fasthttp"
)

func Test_Method(t *testing.T) {
	fi, _ := os.Create("index.html")
	fi.WriteString(`<!DOCTYPE html>

    <meta charset="UTF-8">
    <title>Document</title>

hello static

  `)
	fi.Close()
	app := lu.New()
	fs := DefaultFS
	Static := New(*fs)
	app.Use("/", Static)

	postFallThroughMw := false
	postFallThroughErr := false

	app.Use("/", func(ctx *fasthttp.RequestCtx, next func(error)) {
		postFallThroughMw = true
		postFallThroughErr = false
	})

	go app.Listen(":3000")

	postFallThroughMw = false
	postFallThroughErr = false
	code, body, _ := fasthttp.Get(nil, "http://localhost:3000/")
	if code == 200 && strings.Contains(string(body), "hello static") {
		t.Log("OK")
	} else {
		t.Error("ERROR")
	}

	postFallThroughMw = false
	postFallThroughErr = true
	code, body, _ = fasthttp.Get(nil, "http://localhost:3000/test/xxx")
	if code == 200 && postFallThroughMw && !postFallThroughErr {
		t.Log("OK")
	} else {
		t.Error("ERROR")
	}

	postFallThroughMw = false
	postFallThroughErr = true
	code, _, _ = fasthttp.Get(nil, "http://localhost:3000/test/dist")
	if code == 200 && postFallThroughMw && !postFallThroughErr {
		t.Log("OK")
	} else {
		t.Error("ERROR")
	}
	os.Remove("index.html")
}

func Test_FallThrough1(t *testing.T) {
	os.Mkdir("test", os.ModePerm)
	os.Mkdir("test/dist", os.ModePerm)
	fi, _ := os.Create("test/index.html")
	fi.WriteString(`<!DOCTYPE html>

	  <meta charset="UTF-8">
	  <title>Document</title>

hello test

	`)
	fi.Close()

	fi, _ = os.Create("index.html")
	fi.WriteString(`<!DOCTYPE html>

	  <meta charset="UTF-8">
	  <title>Document</title>

hello static

	`)
	fi.Close()

	app := lu.New()
	fs := DefaultFS
	Static := New(*fs)
	app.Use("/", Static)

	var postFallThroughMw, postFallThroughErr bool
	app.Use("/", func(ctx *fasthttp.RequestCtx, next func(error)) {
		postFallThroughMw = true
		postFallThroughErr = false
	})

	go app.Listen(":3010")

	// method
	postFallThroughMw = false
	postFallThroughErr = true
	code, _, _ := fasthttp.Post(nil, "http://localhost:3010/", nil)
	if code == 200 && postFallThroughMw && !postFallThroughErr {
		t.Log("OK")
	} else {
		t.Error("ERROR")
	}

	postFallThroughMw = false
	postFallThroughErr = true
	code, _, _ = fasthttp.Get(nil, "http://localhost:3010/test/xxx")
	if code == 200 && postFallThroughMw && !postFallThroughErr {
		t.Log("OK")
	} else {
		t.Error("ERROR")
	}

	postFallThroughMw = false
	postFallThroughErr = true
	co, body, _ := fasthttp.Get(nil, "http://localhost:3010/test/index.html")
	if co == 200 && strings.Contains(string(body), "hello test") {
		t.Log("OK")
	} else {
		t.Error("ERROR")
	}

	postFallThroughMw = false
	postFallThroughErr = true
	c, _, _ := fasthttp.Get(nil, "http://localhost:3010/test/dist")
	if c == 200 && postFallThroughMw && !postFallThroughErr {
		t.Log("OK")
	} else {
		t.Error("ERROR")
	}
	os.RemoveAll("test")
	os.Remove("index.html")
}

func Test_FallThrough2(t *testing.T) {
	os.Mkdir("test", os.ModePerm)
	os.Mkdir("test/dist", os.ModePerm)
	fi, _ := os.Create("test/index.html")
	fi.WriteString(`<!DOCTYPE html>

    <meta charset="UTF-8">
    <title>Document</title>

hello test

  `)

	fi.Close()
	app := lu.New()
	fs := DefaultFS
	Static := New(*fs)
	app.Use("/", Static)

	var postFallThroughMw, postFallThroughErr bool
	app.Use("/", func(ctx *fasthttp.RequestCtx, next func(error)) {
		postFallThroughMw = true
		postFallThroughErr = false
		ctx.SetStatusCode(404)
	})

	go app.Listen(":3011")

	// method
	postFallThroughMw = false
	postFallThroughErr = true

	code, _, _ := fasthttp.Post(nil, "http://localhost:3011/", nil)
	if code == 404 && postFallThroughMw && !postFallThroughErr {
		t.Log("OK")
	} else {
		t.Error("ERROR")
	}

	code, _, _ = fasthttp.Get(nil, "http://localhost:3011/test/dist")
	if code == 404 {
		t.Log("OK")
	} else {
		t.Error("ERROR")
	}
	os.RemoveAll("test")
}

func Test_FallThrough3(t *testing.T) {
	os.Mkdir("test", os.ModePerm)
	os.Mkdir("test/dist", os.ModePerm)
	fi, _ := os.Create("test/index.html")
	fi.WriteString(`<!DOCTYPE html>

    <meta charset="UTF-8">
    <title>Document</title>

hello test

  `)

	fi.Close()
	app := lu.New()
	fs := DefaultFS
	fs.IndexNames = nil
	Static := New(*fs)
	app.Use("/", Static)

	var postFallThroughMw, postFallThroughErr bool
	app.Use("/", func(ctx *fasthttp.RequestCtx, next func(error)) {
		postFallThroughMw = true
		postFallThroughErr = false
		ctx.SetStatusCode(404)
	})

	go app.Listen(":3012")

	postFallThroughMw = false
	postFallThroughErr = true
	code, _, _ := fasthttp.Get(nil, "http://localhost:3012/test/")
	if code == 404 && postFallThroughMw && !postFallThroughErr {
		t.Log("OK")
	} else {
		t.Error("ERROR")
	}
	os.RemoveAll("test")
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"compress/gzip"
	"crypto/sha256"
	"flag"
	"go/format"
	"io"
	"io/ioutil"
	"log"
	"mime"
	"os"
	"path"
	"path/filepath"
	"runtime"
	"sort"
	"strings"
	"time"
)

type file struct {
	name  string
	data  string
	mime  string
	mtime time.Time
	// size is the result size before compression. If 0, it means the data is uncompressed
	size int64
	hash []byte
}

func skipFile(file string, excludeSlice []string) bool {
	for _, pattern := range excludeSlice {
		matched, err := path.Match(pattern, file)
		if err != nil {
			log.Fatal(err)
		}
		if matched {
			return true
		}
	}
	return false
}

func processDir(c chan [2]string, dir string, parents []string, excludeSlice []string) {
	files, err := ioutil.ReadDir(dir)
	if err != nil {
		log.Fatal(err)
	}
	for _, file := range files {
		if strings.HasPrefix(file.Name(), ".") {
			continue
		}

		n := filepath.Join(dir, file.Name())
		id := append(parents, file.Name())
		if skipFile(path.Join(id...), excludeSlice) {
			continue
		}

		if file.IsDir() {
			processDir(c, n, id, excludeSlice)
		} else {
			c <- [2]string{n, path.Join(id...)}
		}
	}
}

type fileSlice []*file

func (a fileSlice) Len() int           { return len(a) }
func (a fileSlice) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
func (a fileSlice) Less(i, j int) bool { return a[i].name < a[j].name }

func main() {
	var outputFile, packageName, buildTags, excludeList string
	flag.StringVar(&outputFile, "o", "staticfiles.go", "File to write results to.")
	flag.StringVar(&packageName, "package", "", "Package name of the resulting file. Defaults to name of the resulting file directory")
	flag.StringVar(&buildTags, "build-tags", "", "Build tags to write to the file")
	flag.StringVar(&excludeList, "exclude", "", "Comma-separated patterns to exclude. (e.g. *.scss, sass/ etc.)")
	flag.Parse()
	if flag.NArg() != 1 {
		log.Println("Please pass in a directory to process")
		flag.PrintDefaults()
		os.Exit(1)
		return
	}
	if packageName == "" {
		f, err := filepath.Abs(outputFile)
		if err != nil {
			log.Fatal(err)
		}
		packageName = filepath.Base(filepath.Dir(f))
	}
	excludeSlice := strings.Split(excludeList, ",")

	c := make(chan [2]string, 128)
	go func() {
		for _, arg := range flag.Args() {
			processDir(c, arg, []string{}, excludeSlice)
		}
		close(c)
	}()

	// Used to communicate back all the files that were read and compressed.
	results := make(chan fileSlice)
	process := func() {
		var b bytes.Buffer
		var b2 bytes.Buffer
		hash := sha256.New()
		files := make(fileSlice, 0, 32)
		for asset := range c {
			f, err := os.Open(asset[0])
			if err != nil {
				log.Fatal(err)
			}
			stat, err := f.Stat()
			if err != nil {
				log.Fatal(err)
			}
			if _, err := b.ReadFrom(f); err != nil {
				log.Fatal(err)
			}
			f.Close()
			compressedWriter, _ := gzip.NewWriterLevel(&b2, gzip.BestCompression)
			writer := io.MultiWriter(compressedWriter, hash)
			if _, err := writer.Write(b.Bytes()); err != nil {
				log.Fatal(err)
			}
			compressedWriter.Close()
			if b2.Len() < b.Len() {
				files = append(files, &file{
					name:  asset[1],
					data:  b2.String(),
					mime:  mime.TypeByExtension(filepath.Ext(asset[0])),
					mtime: stat.ModTime(),
					size:  stat.Size(),
					hash:  hash.Sum(nil),
				})
			} else {
				files = append(files, &file{
					name:  asset[1],
					data:  b.String(),
					mime:  mime.TypeByExtension(filepath.Ext(asset[0])),
					mtime: stat.ModTime(),
					hash:  hash.Sum(nil),
				})
			}
			b.Reset()
			b2.Reset()
			hash.Reset()
		}
		results <- files
	}
	// Concurrency! Read and compress in parallel, and combine all the slices at the end.
	concurrency := runtime.NumCPU()
	for i := 0; i < concurrency; i++ {
		go process()
	}
	var files fileSlice
	for i := 0; i < concurrency; i++ {
		files = append(files, (<-results)...)
	}

	// Should sort to make sure the output is stable
	sort.Sort(files)

	var b bytes.Buffer
	if err := GenerateTemplate(&b, packageName, files, buildTags); err != nil {
		log.Fatal(err)
	}
	res, err := format.Source(b.Bytes())
	if err != nil {
		log.Fatal(err)
	}
	if err := os.MkdirAll(filepath.Dir(outputFile), 0755); err != nil {
		log.Fatal(err)
	}
	f, err := os.Create(outputFile)
	if err != nil {
		log.Fatal(err)
	}
	if _, err := f.Write(res); err != nil {
		log.Fatal(err)
	}
}
<commit_msg>Fix argument help<commit_after>package main

import (
	"bytes"
	"compress/gzip"
	"crypto/sha256"
	"flag"
	"go/format"
	"io"
	"io/ioutil"
	"log"
	"mime"
	"os"
	"path"
	"path/filepath"
	"runtime"
	"sort"
	"strings"
	"time"
)

type file struct {
	name  string
	data  string
	mime  string
	mtime time.Time
	// size is the result size before compression. If 0, it means the data is uncompressed
	size int64
	hash []byte
}

func skipFile(file string, excludeSlice []string) bool {
	for _, pattern := range excludeSlice {
		matched, err := path.Match(pattern, file)
		if err != nil {
			log.Fatal(err)
		}
		if matched {
			return true
		}
	}
	return false
}

func processDir(c chan [2]string, dir string, parents []string, excludeSlice []string) {
	files, err := ioutil.ReadDir(dir)
	if err != nil {
		log.Fatal(err)
	}
	for _, file := range files {
		if strings.HasPrefix(file.Name(), ".") {
			continue
		}

		n := filepath.Join(dir, file.Name())
		id := append(parents, file.Name())
		if skipFile(path.Join(id...), excludeSlice) {
			continue
		}

		if file.IsDir() {
			processDir(c, n, id, excludeSlice)
		} else {
			c <- [2]string{n, path.Join(id...)}
		}
	}
}

type fileSlice []*file

func (a fileSlice) Len() int           { return len(a) }
func (a fileSlice) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
func (a fileSlice) Less(i, j int) bool { return a[i].name < a[j].name }

func main() {
	var outputFile, packageName, buildTags, excludeList string
	flag.StringVar(&outputFile, "o", "staticfiles.go", "File to write results to.")
	flag.StringVar(&packageName, "package", "", "Package name of the resulting file. Defaults to name of the resulting file directory")
	flag.StringVar(&buildTags, "build-tags", "", "Build tags to write to the file")
	flag.StringVar(&excludeList, "exclude", "", "Comma-separated patterns to exclude. (e.g. '*.scss,templates/')")
	flag.Parse()
	if flag.NArg() != 1 {
		log.Println("Please pass in a directory to process")
		flag.PrintDefaults()
		os.Exit(1)
		return
	}
	if packageName == "" {
		f, err := filepath.Abs(outputFile)
		if err != nil {
			log.Fatal(err)
		}
		packageName = filepath.Base(filepath.Dir(f))
	}
	excludeSlice := strings.Split(excludeList, ",")

	c := make(chan [2]string, 128)
	go func() {
		for _, arg := range flag.Args() {
			processDir(c, arg, []string{}, excludeSlice)
		}
		close(c)
	}()

	// Used to communicate back all the files that were read and compressed.
	results := make(chan fileSlice)
	process := func() {
		var b bytes.Buffer
		var b2 bytes.Buffer
		hash := sha256.New()
		files := make(fileSlice, 0, 32)
		for asset := range c {
			f, err := os.Open(asset[0])
			if err != nil {
				log.Fatal(err)
			}
			stat, err := f.Stat()
			if err != nil {
				log.Fatal(err)
			}
			if _, err := b.ReadFrom(f); err != nil {
				log.Fatal(err)
			}
			f.Close()
			compressedWriter, _ := gzip.NewWriterLevel(&b2, gzip.BestCompression)
			writer := io.MultiWriter(compressedWriter, hash)
			if _, err := writer.Write(b.Bytes()); err != nil {
				log.Fatal(err)
			}
			compressedWriter.Close()
			if b2.Len() < b.Len() {
				files = append(files, &file{
					name:  asset[1],
					data:  b2.String(),
					mime:  mime.TypeByExtension(filepath.Ext(asset[0])),
					mtime: stat.ModTime(),
					size:  stat.Size(),
					hash:  hash.Sum(nil),
				})
			} else {
				files = append(files, &file{
					name:  asset[1],
					data:  b.String(),
					mime:  mime.TypeByExtension(filepath.Ext(asset[0])),
					mtime: stat.ModTime(),
					hash:  hash.Sum(nil),
				})
			}
			b.Reset()
			b2.Reset()
			hash.Reset()
		}
		results <- files
	}
	// Concurrency! Read and compress in parallel, and combine all the slices at the end.
	concurrency := runtime.NumCPU()
	for i := 0; i < concurrency; i++ {
		go process()
	}
	var files fileSlice
	for i := 0; i < concurrency; i++ {
		files = append(files, (<-results)...)
	}

	// Should sort to make sure the output is stable
	sort.Sort(files)

	var b bytes.Buffer
	if err := GenerateTemplate(&b, packageName, files, buildTags); err != nil {
		log.Fatal(err)
	}
	res, err := format.Source(b.Bytes())
	if err != nil {
		log.Fatal(err)
	}
	if err := os.MkdirAll(filepath.Dir(outputFile), 0755); err != nil {
		log.Fatal(err)
	}
	f, err := os.Create(outputFile)
	if err != nil {
		log.Fatal(err)
	}
	if _, err := f.Write(res); err != nil {
		log.Fatal(err)
	}
}
<|endoftext|>


<commit_before>/*
news-maily - Open-Source Newsletter Mail System

Apache License
Version 2.0, January 2004
http://www.apache.org/licenses/

Copyright 2015 Filip Nikolovski

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

// Package storage provides an interface to the database
// it supports "mysql", "postgres" and "sqlite" database servers
package storage
<commit_msg>Update storage doc<commit_after>// Package storage provides an interface to the database
// it supports "mysql", "postgres" and "sqlite" database servers
package storage
<|endoftext|>


<commit_before>package store

import (
	"os"
	"io"
	"bytes"
	"crypto/rand"
)

const (
	nameLen = 12
)

func SavePost(outDir string, post []byte) (string, error) {
	filename := generateFileName()
	f, err := os.Create(outDir + "/" + filename)
	if err != nil {
		return "", err
	}

	defer f.Close()

	out := post[:0]
	for _, b := range post {
		if b < 32 && b != 10 && b != 13 {
			continue
		}
		out = append(out, b)
	}
	_, err = io.Copy(f, bytes.NewReader(out))

	return filename, err
}

func generateFileName() string {
	c := nameLen
	var dictionary string = "0123456789abcdefghijklmnopqrstuvwxyz"
	var bytes = make([]byte, c)
	rand.Read(bytes)
	for k, v := range bytes {
		 bytes[k] = dictionary[v%byte(len(dictionary))]
	}
	return string(bytes)
}
<commit_msg>Rename a var<commit_after>package store

import (
	"os"
	"io"
	"bytes"
	"crypto/rand"
)

const (
	nameLen = 12
)

func SavePost(outDirectory string, post []byte) (string, error) {
	filename := generateFileName()
	f, err := os.Create(outDirectory + "/" + filename)
	if err != nil {
		return "", err
	}

	defer f.Close()

	out := post[:0]
	for _, b := range post {
		if b < 32 && b != 10 && b != 13 {
			continue
		}
		out = append(out, b)
	}
	_, err = io.Copy(f, bytes.NewReader(out))

	return filename, err
}

func generateFileName() string {
	c := nameLen
	var dictionary string = "0123456789abcdefghijklmnopqrstuvwxyz"
	var bytes = make([]byte, c)
	rand.Read(bytes)
	for k, v := range bytes {
		 bytes[k] = dictionary[v%byte(len(dictionary))]
	}
	return string(bytes)
}
<|endoftext|>


<commit_before>package store

import (
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"hash"
	"io"
	"launchpad.net/gobson/bson"
	"launchpad.net/juju/go/charm"
	"launchpad.net/juju/go/log"
	"launchpad.net/mgo"
	"sort"
)

// TODO: Add meta and config information.

var (
	UpdateConflict  = errors.New("charm update already in progress")
	UpdateIsCurrent = errors.New("charm is already up-to-date")
)

const (
	UpdateTimeout = 600e9
)

// The Store type encapsulates communication with MongoDB to
type Store struct {
	session storeSession
}

// Open creates a new session with the store for adding or
// retrieving charms.
func Open(mongoAddr string) (store *Store, err error) {
	log.Printf("Store opened. Connecting to: %s", mongoAddr)
	store = &Store{}
	session, err := mgo.Mongo(mongoAddr)
	if err != nil {
		log.Printf("Error connecting to MongoDB: %v", err)
		return nil, err
	}
	store = &Store{session: storeSession{session}}
	charms := store.session.Charms()
	err = charms.EnsureIndex(mgo.Index{Key: []string{"url", "revision"}, Unique: true})
	if err != nil {
		log.Printf("Error ensuring charms index: %v", err)
		session.Close()
		return nil, err
	}
	// Put the socket we used on EnsureIndex back in the pool.
	session.Refresh()
	return store, nil
}

// Close terminates the session with store.
func (s *Store) Close() {
	s.session.Close()
}

// dropRevisions returns urls with their revisions removed.
func dropRevisions(urls []*charm.URL) []*charm.URL {
	norev := make([]*charm.URL, len(urls))
	for i := range norev {
		norev[i] = urls[i].WithRevision(-1)
	}
	return norev
}

// AddCharm prepares the store to have a single new charm added to it.
// The charm content should be written onto the store via wc, and once
// wc is closed the charm will be made available at all of the provided
// urls.
// The revisionKey parameter should contain the VCS revision identifier
// that represents the current charm content. An error is returned if
// all of the provided urls are already associated to that revision key.
func (s *Store) AddCharm(urls []*charm.URL, revisionKey string) (wc io.WriteCloser, err error) {
	log.Printf("Trying to add charms %v with key %q...", urls, revisionKey)
	urls = dropRevisions(urls)
	session := s.session.Copy()
	lock, err := session.LockUpdates(urls)
	if err != nil {
		session.Close()
		return nil, err
	}
	defer func() {
		if err != nil {
			lock.Unlock()
			session.Close()
		}
	}()

	maxRev := -1
	newKey := false
	charms := session.Charms()
	charm := charmDoc{}
	for i := range urls {
		urlStr := urls[i].String()
		err := charms.Find(bson.D{{"url", urlStr}}).Sort(bson.D{{"revision", -1}}).One(&charm)
		if err == mgo.NotFound {
			log.Printf("Charm %s not yet in the store.", urlStr)
			newKey = true
			continue
		}
		if charm.RevisionKey != revisionKey {
			log.Printf("Charm %s is out of date with revision key %q.", urlStr, revisionKey)
			newKey = true
		}
		if err != nil {
			log.Printf("Unknown error looking for charm %s: %s", urlStr, err)
			return nil, err
		}
		if charm.Revision > maxRev {
			maxRev = charm.Revision
		}
	}
	if !newKey {
		log.Printf("All charms have revision key %q. Nothing to update.", revisionKey)
		return nil, UpdateIsCurrent
	}
	log.Printf("Preparing writer to add charms with revision %d.", maxRev+1)
	return &writer{session, nil, nil, urls, lock, maxRev + 1, revisionKey}, nil
}

type writer struct {
	session  storeSession
	file     *mgo.GridFile
	sha256   hash.Hash
	urls     []*charm.URL
	lock     *updateMutex
	revision int
	revKey   string
}

// Write streams to the database the data from a charm being inserted
// into the store.
func (w *writer) Write(data []byte) (n int, err error) {
	if w.file == nil {
		w.file, err = w.session.CharmFS().Create("")
		if err != nil {
			log.Printf("Failed to create GridFS file: %v", err)
			return 0, err
		}
		w.sha256 = sha256.New()
		log.Printf("Creating GridFS file with id %q...", w.file.Id().(bson.ObjectId).Hex())
	}
	_, err = w.sha256.Write(data)
	if err != nil {
		panic("hash.Hash should never error")
	}
	return w.file.Write(data)
}

// Close finishes the charm writing process and inserts the final metadata.
// After it completes the charm will be available for consumption.
func (w *writer) Close() error {
	defer w.session.Close()
	defer w.lock.Unlock()
	if w.file == nil {
		return nil
	}
	id := w.file.Id()
	err := w.file.Close()
	if err != nil {
		log.Printf("Failed to close GridFS file: %v", err)
		return err
	}
	charms := w.session.Charms()
	sha256 := hex.EncodeToString(w.sha256.Sum(nil))
	for _, url := range w.urls {
		urlStr := url.String()
		charm := charmDoc{
			urlStr,
			w.revision,
			w.revKey,
			sha256,
			id.(bson.ObjectId),
		}
		err := charms.Insert(&charm)
		if err != nil {
			err = maybeConflict(err)
			log.Printf("Failed to insert new revision of charm %s: %v", urlStr, err)
			return err
		}
	}
	return nil
}

type CharmInfo struct {
	Revision int
	Sha256   string
}

// OpenCharm opens for reading via rc the charm currently available at url.
// rc must necessarily be closed after dealing with it or resources
// will leak.
func (s *Store) OpenCharm(url *charm.URL) (rc io.ReadCloser, info *CharmInfo, err error) {
	session := s.session.Copy()

	log.Debugf("Opening charm %s.", url)
	rev := url.Revision
	url = url.WithRevision(-1)

	charms := session.Charms()
	var charm charmDoc
	var qdoc interface{}
	if rev == -1 {
		qdoc = bson.D{{"url", url.String()}}
	} else {
		qdoc = bson.D{{"url", url.String()}, {"revision", rev}}
	}
	err = charms.Find(qdoc).Sort(bson.D{{"revision", -1}}).One(&charm)
	if err != nil {
		log.Printf("Failed to find charm %s: %v", url, err)
		session.Close()
		return nil, nil, err
	}

	file, err := session.CharmFS().OpenId(charm.FileId)
	if err != nil {
		log.Printf("Failed to open GridFS file for charm %s: %v", url, err)
		session.Close()
		return nil, nil, err
	}
	return &reader{session, file}, &CharmInfo{charm.Revision, charm.Sha256}, nil
}

type reader struct {
	session storeSession
	file    *mgo.GridFile
}

// Read consumes data from the opened charm.
func (r *reader) Read(buf []byte) (n int, err error) {
	return r.file.Read(buf)
}

// Close closes the opened charm, and frees associated resources.
func (r *reader) Close() error {
	err := r.file.Close()
	r.session.Close()
	return err
}

// charmDoc represents the document stored in MongoDB for a single charm revision.
type charmDoc struct {
	URL         string
	Revision    int
	RevisionKey string
	Sha256      string
	FileId      bson.ObjectId
}

// storeSession wraps a mgo.Session ands adds a few convenience methods.
type storeSession struct {
	*mgo.Session
}

// Copy copies the storeSession and its underlying mgo session.
func (s storeSession) Copy() storeSession {
	return storeSession{s.Session.Copy()}
}

// Charms returns the mongo collection where charms are stored.
func (s storeSession) Charms() mgo.Collection {
	return s.DB("juju").C("charms")
}

// Charms returns a mgo.GridFS to read and write charms.
func (s storeSession) CharmFS() *mgo.GridFS {
	return s.DB("juju").GridFS("charmfs")
}

// LockUpdates acquires a server-side lock for updating a single charm
// that is supposed to be made available in all of the provided URLs.
// If the lock can't be acquired, an error will be immediately returned.
// The locks will also expire after UpdateTimeout.
func (s storeSession) LockUpdates(urls []*charm.URL) (m *updateMutex, err error) {
	keys := make([]string, len(urls))
	for i := range urls {
		keys[i] = urls[i].String()
	}
	sort.Strings(keys)
	m = &updateMutex{keys, s.DB("juju").C("locks"), bson.Now()}
	err = m.tryLock()
	if err != nil {
		return nil, err
	}
	return m, nil
}

// updateMutex manages the logic around locking, and is used
// via storeSession.LockUpdates.
type updateMutex struct {
	keys  []string
	locks mgo.Collection
	time  bson.Timestamp
}

// Unlock removes previously acquired locks on all m.keys.
func (m *updateMutex) Unlock() {
	log.Debugf("Unlocking charms for future updates: %v", m.keys)
	for i := len(m.keys) - 1; i >= 0; i-- {
		// Using time below ensures only the proper lock is removed.
		// Can't do much about errors here. Locks will expire anyway.
		m.locks.Remove(bson.D{{"_id", m.keys[i]}, {"time", m.time}})
	}
}

// tryLock tries locking m.keys, one at a time, and succeeds only if it
// can lock all of them in order. The keys should be pre-sorted so that
// two-way conflicts can't happen. In case any of the keys fail to be
// locked, and expiring the old lock doesn't work, tryLock undoes all
// previous locks, and aborts with an error.
func (m *updateMutex) tryLock() error {
	for i, key := range m.keys {
		log.Debugf("Trying to lock charm %s for updates...", key)
		doc := bson.D{{"_id", key}, {"time", m.time}}
		err := m.locks.Insert(doc)
		if err == nil {
			log.Debugf("Charm %s is now locked for updates.", key)
			continue
		}
		if lerr, ok := err.(*mgo.LastError); ok && lerr.Code == 11000 {
			log.Debugf("Charm %s is locked. Trying to expire lock.", key)
			m.tryExpire(key)
			err = m.locks.Insert(doc)
			if err == nil {
				log.Debugf("Charm %s is now locked for updates.", key)
				continue
			}
		}
		// Couldn't lock everyone. Undo previous locks.
		for j := i - 1; j >= 0; j-- {
			// Using time below should be unnecessary, but it's an extra check.
			// Can't do anything about errors here. Lock will expire anyway.
			m.locks.Remove(bson.D{{"_id", m.keys[j]}, {"time", m.time}})
		}
		err = maybeConflict(err)
		log.Printf("Can't lock charms %v for updating: %v", m.keys, err)
		return err
	}
	return nil
}

// tryExpire attempts to remove outdated locks from the database.
func (m *updateMutex) tryExpire(key string) {
	// Ignore errors. If nothing happens the key will continue locked.
	m.locks.Remove(bson.D{{"_id", key}, {"time", bson.D{{"$lt", bson.Now() - UpdateTimeout}}}})
}

// maybeConflict returns an UpdateConflict in case err is a mgo
// insert conflict LastError, or err itself otherwise.
func maybeConflict(err error) error {
	if lerr, ok := err.(*mgo.LastError); ok && lerr.Code == 11000 {
		return UpdateConflict
	}
	return err
}
<commit_msg>Updated documentation.<commit_after>
// The store package is capable of storing and updating charms in a MongoDB
// database, as well as maintaining further information about them such as
// the VCS revision the charm was loaded from and the urls in which the
// charm should be available.
//
// The following MongoDB collections are currently used:
//
//     juju.charms    - Information about the stored charms
//     juju.charmfs.* - GridFS with the charm files
//     juju.locks     - Has unique keys with url of updating charms
//
package store

import (
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"hash"
	"io"
	"launchpad.net/gobson/bson"
	"launchpad.net/juju/go/charm"
	"launchpad.net/juju/go/log"
	"launchpad.net/mgo"
	"sort"
)

// TODO: Add meta and config information.

var (
	UpdateConflict  = errors.New("charm update already in progress")
	UpdateIsCurrent = errors.New("charm is already up-to-date")
)

const (
	UpdateTimeout = 600e9
)

// The Store type encapsulates communication with MongoDB to
type Store struct {
	session storeSession
}

// Open creates a new session with the store for adding or
// retrieving charms.
func Open(mongoAddr string) (store *Store, err error) {
	log.Printf("Store opened. Connecting to: %s", mongoAddr)
	store = &Store{}
	session, err := mgo.Mongo(mongoAddr)
	if err != nil {
		log.Printf("Error connecting to MongoDB: %v", err)
		return nil, err
	}
	store = &Store{session: storeSession{session}}
	charms := store.session.Charms()
	err = charms.EnsureIndex(mgo.Index{Key: []string{"url", "revision"}, Unique: true})
	if err != nil {
		log.Printf("Error ensuring charms index: %v", err)
		session.Close()
		return nil, err
	}
	// Put the socket we used on EnsureIndex back in the pool.
	session.Refresh()
	return store, nil
}

// Close terminates the session with store.
func (s *Store) Close() {
	s.session.Close()
}

// dropRevisions returns urls with their revisions removed.
func dropRevisions(urls []*charm.URL) []*charm.URL {
	norev := make([]*charm.URL, len(urls))
	for i := range norev {
		norev[i] = urls[i].WithRevision(-1)
	}
	return norev
}

// AddCharm prepares the store to have a single new charm added to it.
// The revisionKey parameter should contain the VCS revision identifier
// that represents the current charm content. An error is returned if
// all of the provided urls are already associated to that revision key.
// On success, wc must be used to stream the charm content onto the
// store, and once wc is closed successfully the content will be
// available at all of the provided urls. The returned revision will be
// assigned to the charm, and it's equal to the maximum current revision
// from all of the urls plus one.
func (s *Store) AddCharm(urls []*charm.URL, revisionKey string) (wc io.WriteCloser, revision int, err error) {
	log.Printf("Trying to add charms %v with key %q...", urls, revisionKey)
	urls = dropRevisions(urls)
	session := s.session.Copy()
	lock, err := session.LockUpdates(urls)
	if err != nil {
		session.Close()
		return nil, err
	}
	defer func() {
		if err != nil {
			lock.Unlock()
			session.Close()
		}
	}()

	maxRev := -1
	newKey := false
	charms := session.Charms()
	charm := charmDoc{}
	for i := range urls {
		urlStr := urls[i].String()
		err := charms.Find(bson.D{{"url", urlStr}}).Sort(bson.D{{"revision", -1}}).One(&charm)
		if err == mgo.NotFound {
			log.Printf("Charm %s not yet in the store.", urlStr)
			newKey = true
			continue
		}
		if charm.RevisionKey != revisionKey {
			log.Printf("Charm %s is out of date with revision key %q.", urlStr, revisionKey)
			newKey = true
		}
		if err != nil {
			log.Printf("Unknown error looking for charm %s: %s", urlStr, err)
			return nil, err
		}
		if charm.Revision > maxRev {
			maxRev = charm.Revision
		}
	}
	if !newKey {
		log.Printf("All charms have revision key %q. Nothing to update.", revisionKey)
		return nil, UpdateIsCurrent
	}
	log.Printf("Preparing writer to add charms with revision %d.", maxRev+1)
	return &writer{session, nil, nil, urls, lock, maxRev + 1, revisionKey}, nil
}

type writer struct {
	session  storeSession
	file     *mgo.GridFile
	sha256   hash.Hash
	urls     []*charm.URL
	lock     *updateMutex
	revision int
	revKey   string
}

// Write streams to the database the data from a charm being inserted
// into the store.
func (w *writer) Write(data []byte) (n int, err error) {
	if w.file == nil {
		w.file, err = w.session.CharmFS().Create("")
		if err != nil {
			log.Printf("Failed to create GridFS file: %v", err)
			return 0, err
		}
		w.sha256 = sha256.New()
		log.Printf("Creating GridFS file with id %q...", w.file.Id().(bson.ObjectId).Hex())
	}
	_, err = w.sha256.Write(data)
	if err != nil {
		panic("hash.Hash should never error")
	}
	return w.file.Write(data)
}

// Close finishes the charm writing process and inserts the final metadata.
// After it completes the charm will be available for consumption.
func (w *writer) Close() error {
	defer w.session.Close()
	defer w.lock.Unlock()
	if w.file == nil {
		return nil
	}
	id := w.file.Id()
	err := w.file.Close()
	if err != nil {
		log.Printf("Failed to close GridFS file: %v", err)
		return err
	}
	charms := w.session.Charms()
	sha256 := hex.EncodeToString(w.sha256.Sum(nil))
	for _, url := range w.urls {
		urlStr := url.String()
		charm := charmDoc{
			urlStr,
			w.revision,
			w.revKey,
			sha256,
			id.(bson.ObjectId),
		}
		err := charms.Insert(&charm)
		if err != nil {
			err = maybeConflict(err)
			log.Printf("Failed to insert new revision of charm %s: %v", urlStr, err)
			return err
		}
	}
	return nil
}

type CharmInfo struct {
	Revision int
	Sha256   string
}

// OpenCharm opens for reading via rc the charm currently available at url.
// rc must necessarily be closed after dealing with it or resources
// will leak.
func (s *Store) OpenCharm(url *charm.URL) (rc io.ReadCloser, info *CharmInfo, err error) {
	session := s.session.Copy()

	log.Debugf("Opening charm %s.", url)
	rev := url.Revision
	url = url.WithRevision(-1)

	charms := session.Charms()
	var charm charmDoc
	var qdoc interface{}
	if rev == -1 {
		qdoc = bson.D{{"url", url.String()}}
	} else {
		qdoc = bson.D{{"url", url.String()}, {"revision", rev}}
	}
	err = charms.Find(qdoc).Sort(bson.D{{"revision", -1}}).One(&charm)
	if err != nil {
		log.Printf("Failed to find charm %s: %v", url, err)
		session.Close()
		return nil, nil, err
	}

	file, err := session.CharmFS().OpenId(charm.FileId)
	if err != nil {
		log.Printf("Failed to open GridFS file for charm %s: %v", url, err)
		session.Close()
		return nil, nil, err
	}
	return &reader{session, file}, &CharmInfo{charm.Revision, charm.Sha256}, nil
}

type reader struct {
	session storeSession
	file    *mgo.GridFile
}

// Read consumes data from the opened charm.
func (r *reader) Read(buf []byte) (n int, err error) {
	return r.file.Read(buf)
}

// Close closes the opened charm, and frees associated resources.
func (r *reader) Close() error {
	err := r.file.Close()
	r.session.Close()
	return err
}

// charmDoc represents the document stored in MongoDB for a single charm revision.
type charmDoc struct {
	URL         string
	Revision    int
	RevisionKey string
	Sha256      string
	FileId      bson.ObjectId
}

// storeSession wraps a mgo.Session ands adds a few convenience methods.
type storeSession struct {
	*mgo.Session
}

// Copy copies the storeSession and its underlying mgo session.
func (s storeSession) Copy() storeSession {
	return storeSession{s.Session.Copy()}
}

// Charms returns the mongo collection where charms are stored.
func (s storeSession) Charms() mgo.Collection {
	return s.DB("juju").C("charms")
}

// Charms returns a mgo.GridFS to read and write charms.
func (s storeSession) CharmFS() *mgo.GridFS {
	return s.DB("juju").GridFS("charmfs")
}

// LockUpdates acquires a server-side lock for updating a single charm
// that is supposed to be made available in all of the provided urls.
// If the lock can't be acquired in any of the urls, an error will be
// immediately returned.
// In the usual case, any locking done is undone when an error happens,
// or when m.Unlock is called. In case something else goes wrong, the
// locks will also expire after the period defined in UpdateTimeout.
func (s storeSession) LockUpdates(urls []*charm.URL) (m *updateMutex, err error) {
	keys := make([]string, len(urls))
	for i := range urls {
		keys[i] = urls[i].String()
	}
	sort.Strings(keys)
	m = &updateMutex{keys, s.DB("juju").C("locks"), bson.Now()}
	err = m.tryLock()
	if err != nil {
		return nil, err
	}
	return m, nil
}

// updateMutex manages the logic around locking, and is used
// via storeSession.LockUpdates.
type updateMutex struct {
	keys  []string
	locks mgo.Collection
	time  bson.Timestamp
}

// Unlock removes previously acquired locks on all m.keys.
func (m *updateMutex) Unlock() {
	log.Debugf("Unlocking charms for future updates: %v", m.keys)
	for i := len(m.keys) - 1; i >= 0; i-- {
		// Using time below ensures only the proper lock is removed.
		// Can't do much about errors here. Locks will expire anyway.
		m.locks.Remove(bson.D{{"_id", m.keys[i]}, {"time", m.time}})
	}
}

// tryLock tries locking m.keys, one at a time, and succeeds only if it
// can lock all of them in order. The keys should be pre-sorted so that
// two-way conflicts can't happen. In case any of the keys fail to be
// locked, and expiring the old lock doesn't work, tryLock undoes all
// previous locks, and aborts with an error.
func (m *updateMutex) tryLock() error {
	for i, key := range m.keys {
		log.Debugf("Trying to lock charm %s for updates...", key)
		doc := bson.D{{"_id", key}, {"time", m.time}}
		err := m.locks.Insert(doc)
		if err == nil {
			log.Debugf("Charm %s is now locked for updates.", key)
			continue
		}
		if lerr, ok := err.(*mgo.LastError); ok && lerr.Code == 11000 {
			log.Debugf("Charm %s is locked. Trying to expire lock.", key)
			m.tryExpire(key)
			err = m.locks.Insert(doc)
			if err == nil {
				log.Debugf("Charm %s is now locked for updates.", key)
				continue
			}
		}
		// Couldn't lock everyone. Undo previous locks.
		for j := i - 1; j >= 0; j-- {
			// Using time below should be unnecessary, but it's an extra check.
			// Can't do anything about errors here. Lock will expire anyway.
			m.locks.Remove(bson.D{{"_id", m.keys[j]}, {"time", m.time}})
		}
		err = maybeConflict(err)
		log.Printf("Can't lock charms %v for updating: %v", m.keys, err)
		return err
	}
	return nil
}

// tryExpire attempts to remove outdated locks from the database.
func (m *updateMutex) tryExpire(key string) {
	// Ignore errors. If nothing happens the key will continue locked.
	m.locks.Remove(bson.D{{"_id", key}, {"time", bson.D{{"$lt", bson.Now() - UpdateTimeout}}}})
}

// maybeConflict returns an UpdateConflict in case err is a mgo
// insert conflict LastError, or err itself otherwise.
func maybeConflict(err error) error {
	if lerr, ok := err.(*mgo.LastError); ok && lerr.Code == 11000 {
		return UpdateConflict
	}
	return err
}
<|endoftext|>


<commit_before>package gcache

import (
	"bytes"
	"errors"
	"math/rand"
	"net/http"
	"time"

	"google.golang.org/api/drive/v3"
	"google.golang.org/api/googleapi"
)

var (
	random *rand.Rand
)

func init() {
	random = rand.New(rand.NewSource(time.Now().UnixNano()))
}

// StoreGDrive stores a file that is a given filename on Google Drive.
func StoreGDrive(
	r *http.Request,
	filename string,
	file *drive.File,
	payload *[]byte,
) (
	*drive.File,
	error,
) {

	if filename == "" {
		return nil, errors.New("`filename` must be enough")
	}

	existFile, service, err := getGDriveFile(r, filename, "")
	switch err.(type) {
	case nil:
	case *DriveFileDoesNotExistError:
		file.Name = filename
		file.MimeType = MimeGSuiteDoc
		if file.Parents == nil {
			folderID, err := getRootFolderID(r)
			if err != nil {
				return nil, err
			}
			file.Parents = []string{folderID}
		}
	default:
		return nil, err
	}

	contentType := googleapi.ContentType(mimeTxt)
	var (
		newFile    *drive.File
		clearToken bool
	)

	n := 1
retry:
	payloadReader := bytes.NewReader(*payload)
	<-tokenBucketGDriveAPI
	if existFile == nil {
		newFile, err = service.Files.Create(file).Media(payloadReader, contentType).Do()
	} else {
		newFile, err = service.Files.Update(existFile.Id, file).Media(payloadReader, contentType).Do()
	}

	if err == nil {
		return newFile, nil
	}
	clearToken, n, err = Triable(n, err)
	if err != nil {
		return nil, err
	}
	if clearToken {
		service, err = GetGDriveService(r)
		if err != nil {
			return nil, err
		}
	}
	goto retry
}
<commit_msg>add StoreGDriveByID<commit_after>package gcache

import (
	"bytes"
	"errors"
	"net/http"

	"google.golang.org/api/drive/v3"
	"google.golang.org/api/googleapi"
)

// StoreGDriveByID stores a file that is a given id on Google Drive.
func StoreGDriveByID(
	r *http.Request,
	id string,
	file *drive.File,
	payload *[]byte,
) (
	*drive.File,
	error,
) {

	if id == "" {
		return nil, errors.New("`id` must be enough")
	}

	contentType := googleapi.ContentType(mimeTxt)
	var (
		newFile    *drive.File
		clearToken bool
	)

	n := 1
refresh:
	service, err := GetGDriveService(r)
	if err != nil {
		return nil, err
	}

retry:
	payloadReader := bytes.NewReader(*payload)
	<-tokenBucketGDriveAPI
	newFile, err = service.Files.Update(id, file).Media(payloadReader, contentType).Do()

	if err == nil {
		return newFile, nil
	}
	clearToken, n, err = Triable(n, err)
	if err != nil {
		return nil, err
	}
	if clearToken {
		goto refresh
	}
	goto retry
}
<|endoftext|>


<commit_before>package main

import (
    "bufio"
    "bytes"
    "encoding/json"
    "flag"
    "fmt"
    "io/ioutil"
    "log"
    "net"
    "net/http"
    "os"
    "os/signal"
    "strings"
    "sync"
    "syscall"
    "time"
    "tokenbucket"
)

const _VERSION string = "0.2.0"

type Route struct {
    Bandwidth int64 `json:"bandwidth"`    // Bits per second
    Buffersize uint64 `json:"buffersize"` // Bytes
    Inspect bool `json:"inspect"`         // True = proxy, false = reverse proxy
    Src string `json:"src"`
    Dst []string `json:"dst"`
}

type Itinerary struct {
    Map map[string]Route
}

func sigHandler(ch *chan os.Signal) {
    sig := <-*ch
    fmt.Println("Captured sig", sig)
    os.Exit(3)
}

func main() {
    sigs := make(chan os.Signal, 1)
    signal.Notify(sigs,
                  syscall.SIGHUP,
                  syscall.SIGINT,
                  syscall.SIGQUIT,
                  syscall.SIGABRT,
                  syscall.SIGKILL,
                  syscall.SIGSEGV,
                  syscall.SIGTERM)

    log.SetFlags(log.Ldate | log.Ltime | log.Lshortfile)

    itineraryFile := flag.String("itinerary", "itinerary.json", "file containing source and destination routes")
    flag.Usage = func() {
        fmt.Fprintf(os.Stderr, "version %s\n", _VERSION)
        fmt.Fprintln(os.Stderr, "usage:")
        flag.PrintDefaults()
    }
    flag.Parse()

    itinerary := loadItinerary(itineraryFile)

    for _, m := range itinerary.Map {
        go intercept(m)
    }

    sigHandler(&sigs)
}

func loadItinerary(fileName *string) Itinerary {

    file, _ := os.Open(*fileName)
    decoder := json.NewDecoder(file)
    itinerary := Itinerary{}
    err := decoder.Decode(&itinerary)
    if err != nil {
        log.Println(err.Error())
    }

    return itinerary
}

func intercept(route Route) {
    listener, err := net.Listen("tcp", route.Src)
    if err == nil {
        log.Println("Listening on", route.Src)
        i := 0
        count := 0

        for {
            con, err := listener.Accept()
            if err == nil {
                count = count + 1

                if route.Inspect {
                    go findHttpRoute(count, con, route.Bandwidth / 8, route.Buffersize)
                } else {
                    // Round-robin for now
                    go findRoute(count, con, route.Dst[i], route.Bandwidth / 8, route.Buffersize)
                    i = (i + 1) % len(route.Dst)
                }
            } else {
                log.Println(err.Error())
            }
        }
        listener.Close()
    } else {
        log.Println(err.Error())
    }
}

func findHttpRoute(id int, src net.Conn, bandwidth int64, bufferSize uint64) {
    buf := make([]byte, 16384)
    size, err := src.Read(buf) // Assume one read will include entire HTTP header

    if err == nil {
        reader := bufio.NewReader(strings.NewReader(string(buf[0:size])))
        header, _ := http.ReadRequest(reader)
        dstAddr := header.RequestURI

        dst, err := net.Dial("tcp", dstAddr)

        if err == nil {
            body := ""
            rsp := &http.Response {
                Status:        "200 OK",
                StatusCode:    200,
                Proto:         "HTTP/1.1",
                ProtoMajor:    1,
                ProtoMinor:    1,
                Body:          ioutil.NopCloser(bytes.NewBufferString(body)),
                ContentLength: int64(len(body)),
                Request:       nil,
                Header:        make(http.Header, 0),
            }
            buff := bytes.NewBuffer(nil)
            rsp.Write(buff)
            src.Write(buff.Bytes())
            startDetour(id, src, dst, bandwidth, bufferSize)
        } else {
            src.Close()
            log.Println(err.Error())
        }
    } else {
        src.Close()
        log.Println(err.Error())
    }
}

func findRoute(id int, src net.Conn, dstAddr string, bandwidth int64, bufferSize uint64) {
    dst, err := net.Dial("tcp", dstAddr)

    if err == nil {
        startDetour(id, src, dst, bandwidth, bufferSize)
    } else {
        src.Close()
        log.Println(err.Error())
    }
}

func startDetour(id int, src net.Conn, dst net.Conn, bandwidth int64, bufferSize uint64) {
    log.Println("Opening route", id, ":", src.RemoteAddr().String(), "to", dst.RemoteAddr().String())

    var wg sync.WaitGroup
    wg.Add(1)
    go reroute(&wg, src, dst, bandwidth, bufferSize)
    reroute(nil, dst, src, bandwidth, bufferSize)
    wg.Wait()

    log.Println("Closing route", id, ":", src.RemoteAddr().String(), "to", dst.RemoteAddr().String())
}

func reroute(wg *sync.WaitGroup, src net.Conn, dst net.Conn, bandwidth int64, bufferSize uint64) {
    tb := tokenbucket.New(uint64(bandwidth), 10 * uint64(bandwidth))
    buf := make([]byte, bufferSize)
    defer src.Close()
    defer dst.Close()

    for {
        bytes := tb.Remove(bufferSize)
        if bytes < bufferSize {
            tb.Return(bytes)
            time.Sleep(1 * time.Millisecond)
            continue
        }

        size, err := src.Read(buf)
        if err == nil {
            dst.Write(buf[0:size])
            if size < int(bufferSize) {
                tb.Return(bufferSize - uint64(size))
            }
        } else {
            break
        }
    }

    if wg != nil {
        wg.Done()
    }
}
<commit_msg>Expanded proxy support to include PUT and HEAD requests<commit_after>package main

import (
    "bufio"
    "bytes"
    "encoding/json"
    "flag"
    "fmt"
    "io/ioutil"
    "log"
    "net"
    "net/http"
    "net/url"
    "os"
    "os/signal"
    "strings"
    "sync"
    "syscall"
    "time"
    "tokenbucket"
)

const _VERSION string = "0.2.0"

type Route struct {
    Bandwidth int64 `json:"bandwidth"`    // Bits per second
    Buffersize uint64 `json:"buffersize"` // Bytes
    Inspect bool `json:"inspect"`         // True = proxy, false = reverse proxy
    Src string `json:"src"`
    Dst []string `json:"dst"`
}

type Itinerary struct {
    Map map[string]Route
}

func sigHandler(ch *chan os.Signal) {
    sig := <-*ch
    fmt.Println("Captured sig", sig)
    os.Exit(3)
}

func main() {
    sigs := make(chan os.Signal, 1)
    signal.Notify(sigs,
                  syscall.SIGHUP,
                  syscall.SIGINT,
                  syscall.SIGQUIT,
                  syscall.SIGABRT,
                  syscall.SIGKILL,
                  syscall.SIGSEGV,
                  syscall.SIGTERM)

    log.SetFlags(log.Ldate | log.Ltime | log.Lshortfile)

    itineraryFile := flag.String("itinerary", "itinerary.json", "file containing source and destination routes")
    flag.Usage = func() {
        fmt.Fprintf(os.Stderr, "version %s\n", _VERSION)
        fmt.Fprintln(os.Stderr, "usage:")
        flag.PrintDefaults()
    }
    flag.Parse()

    itinerary := loadItinerary(itineraryFile)

    for _, m := range itinerary.Map {
        go intercept(m)
    }

    sigHandler(&sigs)
}

func loadItinerary(fileName *string) Itinerary {

    file, _ := os.Open(*fileName)
    decoder := json.NewDecoder(file)
    itinerary := Itinerary{}
    err := decoder.Decode(&itinerary)
    if err != nil {
        log.Println(err.Error())
    }

    return itinerary
}

func intercept(route Route) {
    listener, err := net.Listen("tcp", route.Src)
    if err == nil {
        log.Println("Listening on", route.Src)
        i := 0
        count := 0

        for {
            con, err := listener.Accept()
            if err == nil {
                count = count + 1

                if route.Inspect { // Proxy mode
                    go findHttpRoute(count, con, route.Bandwidth / 8, route.Buffersize)
                } else { // Load balancer mode
                    // Round-robin for now
                    go findRoute(count, con, route.Dst[i], route.Bandwidth / 8, route.Buffersize)
                    i = (i + 1) % len(route.Dst)
                }
            } else {
                log.Println(err.Error())
            }
        }
        listener.Close()
    } else {
        log.Println(err.Error())
    }
}

func parseHttpRequestUri(header *http.Request) string {
    var uri string
    if header.Method == "HEAD" || header.Method == "PUT" {
        url, err := url.Parse(header.RequestURI)

        if err == nil {
            if strings.ContainsAny(url.Host, ":") {
                uri = url.Host
            } else {
                switch url.Scheme {
                    case "http":
                        uri = url.Host + ":80"
                    case "https":
                        uri = url.Host + ":443"
                    default:
                        uri = url.Host
                }
            }
        } else {
            log.Println(err.Error())
        }
    } else if header.Method == "CONNECT" {
        uri = header.RequestURI
    }

    return uri
}

func findHttpRoute(id int, src net.Conn, bandwidth int64, bufferSize uint64) {
    buf := make([]byte, 65536)
    size, err := src.Read(buf) // Assume one read will include entire HTTP header

    if err == nil {
        reader := bufio.NewReader(strings.NewReader(string(buf[0:size])))
        header, _ := http.ReadRequest(reader)

        dstAddr := parseHttpRequestUri(header)
        dst, err := net.Dial("tcp", dstAddr)

        if err == nil && header.Method == "CONNECT" {
            body := ""
            rsp := &http.Response {
                Status:        "200 OK",
                StatusCode:    200,
                Proto:         "HTTP/1.1",
                ProtoMajor:    1,
                ProtoMinor:    1,
                Body:          ioutil.NopCloser(bytes.NewBufferString(body)),
                ContentLength: int64(len(body)),
                Request:       nil,
                Header:        make(http.Header, 0),
            }
            buff := bytes.NewBuffer(nil)
            rsp.Write(buff)
            src.Write(buff.Bytes())
            startDetour(id, src, dst, bandwidth, bufferSize)
        } else {
            src.Close()
            log.Println(err.Error())
        }
    } else {
        src.Close()
        log.Println(err.Error())
    }
}

func findRoute(id int, src net.Conn, dstAddr string, bandwidth int64, bufferSize uint64) {
    dst, err := net.Dial("tcp", dstAddr)

    if err == nil {
        startDetour(id, src, dst, bandwidth, bufferSize)
    } else {
        src.Close()
        log.Println(err.Error())
    }
}

func startDetour(id int, src net.Conn, dst net.Conn, bandwidth int64, bufferSize uint64) {
    log.Println("Opening route", id, ":", src.RemoteAddr().String(), "to", dst.RemoteAddr().String())

    var wg sync.WaitGroup
    wg.Add(1)
    go reroute(&wg, src, dst, bandwidth, bufferSize)
    reroute(nil, dst, src, bandwidth, bufferSize)
    wg.Wait()

    log.Println("Closing route", id, ":", src.RemoteAddr().String(), "to", dst.RemoteAddr().String())
}

func reroute(wg *sync.WaitGroup, src net.Conn, dst net.Conn, bandwidth int64, bufferSize uint64) {
    tb := tokenbucket.New(uint64(bandwidth), 10 * uint64(bandwidth))
    buf := make([]byte, bufferSize)
    defer src.Close()
    defer dst.Close()

    for {
        bytes := tb.Remove(bufferSize)
        if bytes < bufferSize {
            tb.Return(bytes)
            time.Sleep(1 * time.Millisecond)
            continue
        }

        size, err := src.Read(buf)
        if err == nil {
            dst.Write(buf[0:size])
            if size < int(bufferSize) {
                tb.Return(bufferSize - uint64(size))
            }
        } else {
            break
        }
    }

    if wg != nil {
        wg.Done()
    }
}
<|endoftext|>


<commit_before>package dns

import (
	"crypto/sha1"
	"hash"
	"strings"
)

type saltWireFmt struct {
	Salt string `dns:"size-hex"`
}

// HashName hashes a string (label) according to RFC 5155. It returns the hashed string in uppercase.
func HashName(label string, ha uint8, iter uint16, salt string) string {
	saltwire := new(saltWireFmt)
	saltwire.Salt = salt
	wire := make([]byte, DefaultMsgSize)
	n, err := packSaltWire(saltwire, wire)
	if err != nil {
		return ""
	}
	wire = wire[:n]
	name := make([]byte, 255)
	off, err := PackDomainName(strings.ToLower(label), name, 0, nil, false)
	if err != nil {
		return ""
	}
	name = name[:off]
	var s hash.Hash
	switch ha {
	case SHA1:
		s = sha1.New()
	default:
		return ""
	}

	// k = 0
	s.Write(name)
	s.Write(wire)
	nsec3 := s.Sum(nil)
	// k > 0
	for k := uint16(0); k < iter; k++ {
		s.Reset()
		s.Write(nsec3)
		s.Write(wire)
		nsec3 = s.Sum(nsec3[:0])
	}
	return toBase32(nsec3)
}

// Cover returns true if a name is covered by the NSEC3 record
func (rr *NSEC3) Cover(name string) bool {
	nameHash := HashName(name, rr.Hash, rr.Iterations, rr.Salt)
	owner := strings.ToUpper(rr.Hdr.Name)
	labelIndices := Split(owner)
	if len(labelIndices) < 2 {
		return false
	}
	ownerHash := owner[:labelIndices[1]-1]
	ownerZone := owner[labelIndices[1]:]
	if !IsSubDomain(ownerZone, strings.ToUpper(name)) { // name is outside owner zone
		return false
	}

	nextHash := rr.NextDomain

	// if empty interval found, try cover wildcard hashes so nameHash shouldn't match with ownerHash
	if ownerHash == nextHash && nameHash != ownerHash { // empty interval
		return true
	}
	if ownerHash > nextHash { // end of zone
		if nameHash > ownerHash { // covered since there is nothing after ownerHash
			return true
		}
		return nameHash < nextHash // if nameHash is before beginning of zone it is covered
	}
	if nameHash < ownerHash { // nameHash is before ownerHash, not covered
		return false
	}
	return nameHash < nextHash // if nameHash is before nextHash is it covered (between ownerHash and nextHash)
}

// Match returns true if a name matches the NSEC3 record
func (rr *NSEC3) Match(name string) bool {
	nameHash := HashName(name, rr.Hash, rr.Iterations, rr.Salt)
	owner := strings.ToUpper(rr.Hdr.Name)
	labelIndices := Split(owner)
	if len(labelIndices) < 2 {
		return false
	}
	ownerHash := owner[:labelIndices[1]-1]
	ownerZone := owner[labelIndices[1]:]
	if !IsSubDomain(ownerZone, strings.ToUpper(name)) { // name is outside owner zone
		return false
	}
	if ownerHash == nameHash {
		return true
	}
	return false
}

func packSaltWire(sw *saltWireFmt, msg []byte) (int, error) {
	off, err := packStringHex(sw.Salt, msg, 0)
	if err != nil {
		return off, err
	}
	return off, nil
}
<commit_msg>Remove saltWireFmt from nsecx.go<commit_after>package dns

import (
	"crypto/sha1"
	"hash"
	"strings"
)

// HashName hashes a string (label) according to RFC 5155. It returns the hashed string in uppercase.
func HashName(label string, ha uint8, iter uint16, salt string) string {
	wire := make([]byte, DefaultMsgSize)
	n, err := packStringHex(salt, wire, 0)
	if err != nil {
		return ""
	}
	wire = wire[:n]
	name := make([]byte, 255)
	off, err := PackDomainName(strings.ToLower(label), name, 0, nil, false)
	if err != nil {
		return ""
	}
	name = name[:off]
	var s hash.Hash
	switch ha {
	case SHA1:
		s = sha1.New()
	default:
		return ""
	}

	// k = 0
	s.Write(name)
	s.Write(wire)
	nsec3 := s.Sum(nil)
	// k > 0
	for k := uint16(0); k < iter; k++ {
		s.Reset()
		s.Write(nsec3)
		s.Write(wire)
		nsec3 = s.Sum(nsec3[:0])
	}
	return toBase32(nsec3)
}

// Cover returns true if a name is covered by the NSEC3 record
func (rr *NSEC3) Cover(name string) bool {
	nameHash := HashName(name, rr.Hash, rr.Iterations, rr.Salt)
	owner := strings.ToUpper(rr.Hdr.Name)
	labelIndices := Split(owner)
	if len(labelIndices) < 2 {
		return false
	}
	ownerHash := owner[:labelIndices[1]-1]
	ownerZone := owner[labelIndices[1]:]
	if !IsSubDomain(ownerZone, strings.ToUpper(name)) { // name is outside owner zone
		return false
	}

	nextHash := rr.NextDomain

	// if empty interval found, try cover wildcard hashes so nameHash shouldn't match with ownerHash
	if ownerHash == nextHash && nameHash != ownerHash { // empty interval
		return true
	}
	if ownerHash > nextHash { // end of zone
		if nameHash > ownerHash { // covered since there is nothing after ownerHash
			return true
		}
		return nameHash < nextHash // if nameHash is before beginning of zone it is covered
	}
	if nameHash < ownerHash { // nameHash is before ownerHash, not covered
		return false
	}
	return nameHash < nextHash // if nameHash is before nextHash is it covered (between ownerHash and nextHash)
}

// Match returns true if a name matches the NSEC3 record
func (rr *NSEC3) Match(name string) bool {
	nameHash := HashName(name, rr.Hash, rr.Iterations, rr.Salt)
	owner := strings.ToUpper(rr.Hdr.Name)
	labelIndices := Split(owner)
	if len(labelIndices) < 2 {
		return false
	}
	ownerHash := owner[:labelIndices[1]-1]
	ownerZone := owner[labelIndices[1]:]
	if !IsSubDomain(ownerZone, strings.ToUpper(name)) { // name is outside owner zone
		return false
	}
	if ownerHash == nameHash {
		return true
	}
	return false
}
<|endoftext|>