<commit_before>package pretty

import (
	"sort"
)

// Options is Pretty options
type Options struct {
	// Width is an max column width for single line arrays
	// Default is 80
	Width int
	// Prefix is a prefix for all lines
	// Default is an empty string
	Prefix string
	// Indent is the nested indentation
	// Default is two spaces
	Indent string
	// SortKeys will sort the keys alphabetically
	// Default is false
	SortKeys bool
}

// DefaultOptions is the default options for pretty formats.
var DefaultOptions = &Options{Width: 80, Prefix: "", Indent: "  ", SortKeys: false}

// Pretty converts the input json into a more human readable format where each
// element is on it's own line with clear indentation.
func Pretty(json []byte) []byte { return PrettyOptions(json, nil) }

// PrettyOptions is like Pretty but with customized options.
func PrettyOptions(json []byte, opts *Options) []byte {
	if opts == nil {
		opts = DefaultOptions
	}
	buf := make([]byte, 0, len(json))
	if len(opts.Prefix) != 0 {
		buf = append(buf, opts.Prefix...)
	}
	buf, _, _, _ = appendPrettyAny(buf, json, 0, true,
		opts.Width, opts.Prefix, opts.Indent, opts.SortKeys,
		0, 0, -1)
	return buf
}

// Ugly removes insignificant space characters from the input json byte slice
// and returns the compacted result.
func Ugly(json []byte) []byte {
	buf := make([]byte, 0, len(json))
	return ugly(buf, json)
}

// UglyInPlace removes insignificant space characters from the input json
// byte slice and returns the compacted result. This method reuses the
// input json buffer to avoid allocations. Do not use the original bytes
// slice upon return.
func UglyInPlace(json []byte) []byte { return ugly(json, json) }

func ugly(dst, src []byte) []byte {
	dst = dst[:0]
	for i := 0; i < len(src); i++ {
		if src[i] > ' ' {
			dst = append(dst, src[i])
			if src[i] == '"' {
				for i = i + 1; i < len(src); i++ {
					dst = append(dst, src[i])
					if src[i] == '"' {
						j := i - 1
						for ; ; j-- {
							if src[j] != '\\' {
								break
							}
						}
						if (j-i)%2 != 0 {
							break
						}
					}
				}
			}
		}
	}
	return dst
}

func appendPrettyAny(buf, json []byte, i int, pretty bool, width int, prefix, indent string, sortkeys bool, tabs, nl, max int) ([]byte, int, int, bool) {
	for ; i < len(json); i++ {
		if json[i] <= ' ' {
			continue
		}
		if json[i] == '"' {
			return appendPrettyString(buf, json, i, nl)
		}
		if (json[i] >= '0' && json[i] <= '9') || json[i] == '-' {
			return appendPrettyNumber(buf, json, i, nl)
		}
		if json[i] == '{' {
			return appendPrettyObject(buf, json, i, '{', '}', pretty, width, prefix, indent, sortkeys, tabs, nl, max)
		}
		if json[i] == '[' {
			return appendPrettyObject(buf, json, i, '[', ']', pretty, width, prefix, indent, sortkeys, tabs, nl, max)
		}
		switch json[i] {
		case 't':
			return append(buf, 't', 'r', 'u', 'e'), i + 4, nl, true
		case 'f':
			return append(buf, 'f', 'a', 'l', 's', 'e'), i + 5, nl, true
		case 'n':
			return append(buf, 'n', 'u', 'l', 'l'), i + 4, nl, true
		}
	}
	return buf, i, nl, true
}

type pair struct {
	kstart, kend int
	vstart, vend int
}

type byKey struct {
	sorted bool
	json   []byte
	pairs  []pair
}

func (arr *byKey) Len() int {
	return len(arr.pairs)
}
func (arr *byKey) Less(i, j int) bool {
	key1 := arr.json[arr.pairs[i].kstart+1 : arr.pairs[i].kend-1]
	key2 := arr.json[arr.pairs[j].kstart+1 : arr.pairs[j].kend-1]
	return string(key1) < string(key2)
}
func (arr *byKey) Swap(i, j int) {
	arr.pairs[i], arr.pairs[j] = arr.pairs[j], arr.pairs[i]
	arr.sorted = true
}

func appendPrettyObject(buf, json []byte, i int, open, close byte, pretty bool, width int, prefix, indent string, sortkeys bool, tabs, nl, max int) ([]byte, int, int, bool) {
	var ok bool
	if width > 0 {
		if pretty && open == '[' && max == -1 {
			// here we try to create a single line array
			max := width - (len(buf) - nl)
			if max > 3 {
				s1, s2 := len(buf), i
				buf, i, _, ok = appendPrettyObject(buf, json, i, '[', ']', false, width, prefix, "", sortkeys, 0, 0, max)
				if ok && len(buf)-s1 <= max {
					return buf, i, nl, true
				}
				buf = buf[:s1]
				i = s2
			}
		} else if max != -1 && open == '{' {
			return buf, i, nl, false
		}
	}
	buf = append(buf, open)
	i++
	var pairs []pair
	if open == '{' && sortkeys {
		pairs = make([]pair, 0, 8)
	}
	var n int
	for ; i < len(json); i++ {
		if json[i] <= ' ' {
			continue
		}
		if json[i] == close {
			if pretty {
				if open == '{' && sortkeys {
					buf = sortPairs(json, buf, pairs)
				}
				if n > 0 {
					nl = len(buf)
					buf = append(buf, '\n')
				}
				if buf[len(buf)-1] != open {
					buf = appendTabs(buf, prefix, indent, tabs)
				}
			}
			buf = append(buf, close)
			return buf, i + 1, nl, open != '{'
		}
		if open == '[' || json[i] == '"' {
			if n > 0 {
				buf = append(buf, ',')
				if width != -1 {
					buf = append(buf, ' ')
				}
			}
			var p pair
			if pretty {
				nl = len(buf)
				buf = append(buf, '\n')
				if open == '{' && sortkeys {
					p.kstart = i
					p.vstart = len(buf)
				}
				buf = appendTabs(buf, prefix, indent, tabs+1)
			}
			if open == '{' {
				buf, i, nl, _ = appendPrettyString(buf, json, i, nl)
				if sortkeys {
					p.kend = i
				}
				buf = append(buf, ':')
				if pretty {
					buf = append(buf, ' ')
				}
			}
			buf, i, nl, ok = appendPrettyAny(buf, json, i, pretty, width, prefix, indent, sortkeys, tabs+1, nl, max)
			if max != -1 && !ok {
				return buf, i, nl, false
			}
			if pretty && open == '{' && sortkeys {
				p.vend = len(buf)
				if p.kstart > p.kend || p.vstart > p.vend {
					// bad data. disable sorting
					sortkeys = false
				} else {
					pairs = append(pairs, p)
				}
			}
			i--
			n++
		}
	}
	return buf, i, nl, open != '{'
}
func sortPairs(json, buf []byte, pairs []pair) []byte {
	if len(pairs) == 0 {
		return buf
	}
	vstart := pairs[0].vstart
	vend := pairs[len(pairs)-1].vend
	arr := byKey{false, json, pairs}
	sort.Sort(&arr)
	if !arr.sorted {
		return buf
	}
	nbuf := make([]byte, 0, vend-vstart)
	for i, p := range pairs {
		nbuf = append(nbuf, buf[p.vstart:p.vend]...)
		if i < len(pairs)-1 {
			nbuf = append(nbuf, ',')
			nbuf = append(nbuf, '\n')
		}
	}
	return append(buf[:vstart], nbuf...)
}

func appendPrettyString(buf, json []byte, i, nl int) ([]byte, int, int, bool) {
	s := i
	i++
	for ; i < len(json); i++ {
		if json[i] == '"' {
			var sc int
			for j := i - 1; j > s; j-- {
				if json[j] == '\\' {
					sc++
				} else {
					break
				}
			}
			if sc%2 == 1 {
				continue
			}
			i++
			break
		}
	}
	return append(buf, json[s:i]...), i, nl, true
}

func appendPrettyNumber(buf, json []byte, i, nl int) ([]byte, int, int, bool) {
	s := i
	i++
	for ; i < len(json); i++ {
		if json[i] <= ' ' || json[i] == ',' || json[i] == ':' || json[i] == ']' || json[i] == '}' {
			break
		}
	}
	return append(buf, json[s:i]...), i, nl, true
}

func appendTabs(buf []byte, prefix, indent string, tabs int) []byte {
	if len(prefix) != 0 {
		buf = append(buf, prefix...)
	}
	for i := 0; i < tabs; i++ {
		buf = append(buf, ' ', ' ')
	}
	return buf
}
<commit_msg>fixed custom indents<commit_after>package pretty

import (
	"sort"
)

// Options is Pretty options
type Options struct {
	// Width is an max column width for single line arrays
	// Default is 80
	Width int
	// Prefix is a prefix for all lines
	// Default is an empty string
	Prefix string
	// Indent is the nested indentation
	// Default is two spaces
	Indent string
	// SortKeys will sort the keys alphabetically
	// Default is false
	SortKeys bool
}

// DefaultOptions is the default options for pretty formats.
var DefaultOptions = &Options{Width: 80, Prefix: "", Indent: "  ", SortKeys: false}

// Pretty converts the input json into a more human readable format where each
// element is on it's own line with clear indentation.
func Pretty(json []byte) []byte { return PrettyOptions(json, nil) }

// PrettyOptions is like Pretty but with customized options.
func PrettyOptions(json []byte, opts *Options) []byte {
	if opts == nil {
		opts = DefaultOptions
	}
	buf := make([]byte, 0, len(json))
	if len(opts.Prefix) != 0 {
		buf = append(buf, opts.Prefix...)
	}
	buf, _, _, _ = appendPrettyAny(buf, json, 0, true,
		opts.Width, opts.Prefix, opts.Indent, opts.SortKeys,
		0, 0, -1)
	return buf
}

// Ugly removes insignificant space characters from the input json byte slice
// and returns the compacted result.
func Ugly(json []byte) []byte {
	buf := make([]byte, 0, len(json))
	return ugly(buf, json)
}

// UglyInPlace removes insignificant space characters from the input json
// byte slice and returns the compacted result. This method reuses the
// input json buffer to avoid allocations. Do not use the original bytes
// slice upon return.
func UglyInPlace(json []byte) []byte { return ugly(json, json) }

func ugly(dst, src []byte) []byte {
	dst = dst[:0]
	for i := 0; i < len(src); i++ {
		if src[i] > ' ' {
			dst = append(dst, src[i])
			if src[i] == '"' {
				for i = i + 1; i < len(src); i++ {
					dst = append(dst, src[i])
					if src[i] == '"' {
						j := i - 1
						for ; ; j-- {
							if src[j] != '\\' {
								break
							}
						}
						if (j-i)%2 != 0 {
							break
						}
					}
				}
			}
		}
	}
	return dst
}

func appendPrettyAny(buf, json []byte, i int, pretty bool, width int, prefix, indent string, sortkeys bool, tabs, nl, max int) ([]byte, int, int, bool) {
	for ; i < len(json); i++ {
		if json[i] <= ' ' {
			continue
		}
		if json[i] == '"' {
			return appendPrettyString(buf, json, i, nl)
		}
		if (json[i] >= '0' && json[i] <= '9') || json[i] == '-' {
			return appendPrettyNumber(buf, json, i, nl)
		}
		if json[i] == '{' {
			return appendPrettyObject(buf, json, i, '{', '}', pretty, width, prefix, indent, sortkeys, tabs, nl, max)
		}
		if json[i] == '[' {
			return appendPrettyObject(buf, json, i, '[', ']', pretty, width, prefix, indent, sortkeys, tabs, nl, max)
		}
		switch json[i] {
		case 't':
			return append(buf, 't', 'r', 'u', 'e'), i + 4, nl, true
		case 'f':
			return append(buf, 'f', 'a', 'l', 's', 'e'), i + 5, nl, true
		case 'n':
			return append(buf, 'n', 'u', 'l', 'l'), i + 4, nl, true
		}
	}
	return buf, i, nl, true
}

type pair struct {
	kstart, kend int
	vstart, vend int
}

type byKey struct {
	sorted bool
	json   []byte
	pairs  []pair
}

func (arr *byKey) Len() int {
	return len(arr.pairs)
}
func (arr *byKey) Less(i, j int) bool {
	key1 := arr.json[arr.pairs[i].kstart+1 : arr.pairs[i].kend-1]
	key2 := arr.json[arr.pairs[j].kstart+1 : arr.pairs[j].kend-1]
	return string(key1) < string(key2)
}
func (arr *byKey) Swap(i, j int) {
	arr.pairs[i], arr.pairs[j] = arr.pairs[j], arr.pairs[i]
	arr.sorted = true
}

func appendPrettyObject(buf, json []byte, i int, open, close byte, pretty bool, width int, prefix, indent string, sortkeys bool, tabs, nl, max int) ([]byte, int, int, bool) {
	var ok bool
	if width > 0 {
		if pretty && open == '[' && max == -1 {
			// here we try to create a single line array
			max := width - (len(buf) - nl)
			if max > 3 {
				s1, s2 := len(buf), i
				buf, i, _, ok = appendPrettyObject(buf, json, i, '[', ']', false, width, prefix, "", sortkeys, 0, 0, max)
				if ok && len(buf)-s1 <= max {
					return buf, i, nl, true
				}
				buf = buf[:s1]
				i = s2
			}
		} else if max != -1 && open == '{' {
			return buf, i, nl, false
		}
	}
	buf = append(buf, open)
	i++
	var pairs []pair
	if open == '{' && sortkeys {
		pairs = make([]pair, 0, 8)
	}
	var n int
	for ; i < len(json); i++ {
		if json[i] <= ' ' {
			continue
		}
		if json[i] == close {
			if pretty {
				if open == '{' && sortkeys {
					buf = sortPairs(json, buf, pairs)
				}
				if n > 0 {
					nl = len(buf)
					buf = append(buf, '\n')
				}
				if buf[len(buf)-1] != open {
					buf = appendTabs(buf, prefix, indent, tabs)
				}
			}
			buf = append(buf, close)
			return buf, i + 1, nl, open != '{'
		}
		if open == '[' || json[i] == '"' {
			if n > 0 {
				buf = append(buf, ',')
				if width != -1 {
					buf = append(buf, ' ')
				}
			}
			var p pair
			if pretty {
				nl = len(buf)
				buf = append(buf, '\n')
				if open == '{' && sortkeys {
					p.kstart = i
					p.vstart = len(buf)
				}
				buf = appendTabs(buf, prefix, indent, tabs+1)
			}
			if open == '{' {
				buf, i, nl, _ = appendPrettyString(buf, json, i, nl)
				if sortkeys {
					p.kend = i
				}
				buf = append(buf, ':')
				if pretty {
					buf = append(buf, ' ')
				}
			}
			buf, i, nl, ok = appendPrettyAny(buf, json, i, pretty, width, prefix, indent, sortkeys, tabs+1, nl, max)
			if max != -1 && !ok {
				return buf, i, nl, false
			}
			if pretty && open == '{' && sortkeys {
				p.vend = len(buf)
				if p.kstart > p.kend || p.vstart > p.vend {
					// bad data. disable sorting
					sortkeys = false
				} else {
					pairs = append(pairs, p)
				}
			}
			i--
			n++
		}
	}
	return buf, i, nl, open != '{'
}
func sortPairs(json, buf []byte, pairs []pair) []byte {
	if len(pairs) == 0 {
		return buf
	}
	vstart := pairs[0].vstart
	vend := pairs[len(pairs)-1].vend
	arr := byKey{false, json, pairs}
	sort.Sort(&arr)
	if !arr.sorted {
		return buf
	}
	nbuf := make([]byte, 0, vend-vstart)
	for i, p := range pairs {
		nbuf = append(nbuf, buf[p.vstart:p.vend]...)
		if i < len(pairs)-1 {
			nbuf = append(nbuf, ',')
			nbuf = append(nbuf, '\n')
		}
	}
	return append(buf[:vstart], nbuf...)
}

func appendPrettyString(buf, json []byte, i, nl int) ([]byte, int, int, bool) {
	s := i
	i++
	for ; i < len(json); i++ {
		if json[i] == '"' {
			var sc int
			for j := i - 1; j > s; j-- {
				if json[j] == '\\' {
					sc++
				} else {
					break
				}
			}
			if sc%2 == 1 {
				continue
			}
			i++
			break
		}
	}
	return append(buf, json[s:i]...), i, nl, true
}

func appendPrettyNumber(buf, json []byte, i, nl int) ([]byte, int, int, bool) {
	s := i
	i++
	for ; i < len(json); i++ {
		if json[i] <= ' ' || json[i] == ',' || json[i] == ':' || json[i] == ']' || json[i] == '}' {
			break
		}
	}
	return append(buf, json[s:i]...), i, nl, true
}

func appendTabs(buf []byte, prefix, indent string, tabs int) []byte {
	if len(prefix) != 0 {
		buf = append(buf, prefix...)
	}
	if len(indent) == 2 && indent[0] == ' ' && indent[1] == ' ' {
		for i := 0; i < tabs; i++ {
			buf = append(buf, ' ', ' ')
		}
	} else {
		for i := 0; i < tabs; i++ {
			buf = append(buf, indent...)
		}
	}
	return buf
}
<|endoftext|>


<commit_before>package server

import (
	"context"
	"fmt"

	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"

	log "github.com/Sirupsen/logrus"
	"github.com/agave/ah-microservices/services/users/db"
	"github.com/agave/ah-microservices/services/users/user"
	userGen "github.com/agave/ah-microservices/services/users/user/generated"
	"github.com/agave/ah-microservices/services/users/util"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/suite"
)

type ServerFunctionalSuite struct {
	suite.Suite
	A *assert.Assertions
}

func (s *ServerFunctionalSuite) SetupSuite() {
	var err error
	s.A = assert.New(s.T())

	// db
	util.Config.DB.DBName = util.Config.DB.DBName + "_test"

	url := util.FormatDBURL("postgres", &util.Config.DB)
	err = db.InitDB("postgres", url)
	if err != nil {
		log.WithFields(log.Fields{
			"error": err, "url": url,
		}).Fatalln("Error starting database")
		return
	}

	err = db.Engine.CreateTables(&user.Users{})
	if err != nil {
		log.WithField("error", err).Fatalln("Error creating tables")
	}
}

func (s *ServerFunctionalSuite) TearDownSuite() {
	if ok, _ := db.Engine.IsTableExist(&user.Users{}); ok {
		db.Engine.DropTables(&user.Users{})
	}
	if ok, _ := db.Engine.IsTableExist(&user.HeldBalance{}); ok {
		db.Engine.DropTables(&user.HeldBalance{})
	}
	db.Engine.Close()
}

func (s *ServerFunctionalSuite) TestGetUser() {
	db.Engine.CreateTables(user.Users{})
	type testCases struct {
		Ins  []*userGen.Id
		Outs []codes.Code
	}

	h := &userGen.Create{Email: "admin@admin.com", Balance: 0}
	p, _ := grpcClient.CreateUser(context.TODO(), h)

	tc := testCases{
		Ins: []*userGen.Id{
			&userGen.Id{Guid: "1", Id: p.GetId()}, // happy
			&userGen.Id{Guid: "2", Id: "10000"},   // not found
			&userGen.Id{Guid: "3", Id: "-99999"},
		},
		Outs: []codes.Code{
			codes.OK,
			codes.NotFound,
			codes.NotFound,
		},
	}

	for i, v := range tc.Ins {
		p, err := grpcClient.GetUser(context.TODO(), v)
		s.A.Equal(tc.Outs[i], grpc.Code(err), "GetUser should return expected error")
		if tc.Outs[i] == codes.OK {
			s.A.Equal(v.GetId(), p.GetId(), "Should return same id")
			s.A.NotEmpty(p.GetEmail(), "Email shouldn't be empty")
		}
	}

	//trigger db error
	db.Engine.DropTables(&user.Users{})
	p, err := grpcClient.GetUser(context.TODO(), &userGen.Id{Id: "1"})
	s.A.Nil(p)
	s.A.Equal(codes.Internal, grpc.Code(err))
}

func (s *ServerFunctionalSuite) TestCreateUser() {
	db.Engine.CreateTables(user.Users{})
	type testCases struct {
		Ins  []*userGen.Create
		Outs []codes.Code
	}

	tc := testCases{
		Ins: []*userGen.Create{
			&userGen.Create{Email: "admin@admin.com", Balance: 0},
			&userGen.Create{Email: "admin@admin.com", Balance: 0},
			&userGen.Create{Email: "admin2@admin.com", Balance: -42},
			&userGen.Create{Email: "admin3@admin.com", Balance: 42},
			&userGen.Create{Email: "admin4@admin.com", Balance: 3456.3453},
			&userGen.Create{Email: "", Balance: 3456.3453},
		},
		Outs: []codes.Code{
			codes.OK,
			codes.AlreadyExists,
			codes.InvalidArgument,
			codes.OK,
			codes.OK,
			codes.InvalidArgument,
		},
	}

	for i, c := range tc.Ins {
		p, err := grpcClient.CreateUser(context.TODO(), c)
		s.A.Equal(tc.Outs[i], grpc.Code(err), "Create should return expected error")
		if tc.Outs[i] == codes.OK {
			s.A.NotNil(p)
			s.A.Equal(c.GetEmail(), p.GetEmail(), "Insert should return same Email")
			s.A.Equal(c.GetBalance(), p.GetBalance(), "Insert should return same Balance")
		}
	}

	//trigger db error
	db.Engine.DropTables(&user.Users{})
	crt := &userGen.Create{
		Email:   "not empty",
		Balance: 100.00,
	}
	p, err := grpcClient.CreateUser(context.TODO(), crt)
	s.A.Nil(p)
	s.A.Equal(codes.Unknown, grpc.Code(err))
}

func (s *ServerFunctionalSuite) TestUpdateUser() {
	db.Engine.CreateTables(user.Users{})
	//insert fixtures
	createFixture := &userGen.Create{
		Email:   "user@domain.org",
		Balance: 10.00,
	}
	u, err := grpcClient.CreateUser(context.TODO(), createFixture)
	s.A.Nil(err)
	s.A.Equal(createFixture.GetEmail(), u.GetEmail(), "Should return inserted fixture's email")
	s.A.Equal(createFixture.GetBalance(), u.GetBalance(), "Should return inserted fixture's balance")
	s.A.NotEmpty(u.GetId(), "Should return with a valid Id")

	//change email
	updateParams := &userGen.UpdateUserParams{
		Id:            u.GetId(),
		Email:         "user@domain.com",
		Balance:       0.0,
		ChangeBalance: false,
	}

	nu, err := grpcClient.UpdateUser(context.TODO(), updateParams)
	s.A.Nil(err)
	s.A.Equal(updateParams.GetEmail(), nu.GetEmail(), "User email should be updated when not empty")
	s.A.NotEqual(updateParams.GetBalance(), nu.GetBalance(), "Balance shouldn't be updated when ChangeBalanceFlag is false")
	s.A.Equal(updateParams.GetId(), nu.GetId(), "Id never changes")

	// change balance
	updateParams = &userGen.UpdateUserParams{
		Id:            u.GetId(),
		Email:         "", //email remains the same
		Balance:       1000.0,
		ChangeBalance: true,
	}

	nu, err = grpcClient.UpdateUser(context.TODO(), updateParams)
	s.A.Nil(err)
	s.A.NotEqual(updateParams.GetEmail(), nu.GetEmail(), "User email should not update when empty")
	s.A.Equal(updateParams.GetBalance(), nu.GetBalance(), "Balance should be updated")

	// invalid/non existent id
	updateParams = &userGen.UpdateUserParams{
		Id:            "-666",
		Email:         "", //email remains the same
		Balance:       1000.0,
		ChangeBalance: true,
	}

	nu, err = grpcClient.UpdateUser(context.TODO(), updateParams)
	s.A.NotNil(err)
	s.A.Equal(codes.NotFound, grpc.Code(err))
	s.A.Nil(nu)

	//no changes
	updateParams = &userGen.UpdateUserParams{
		Id:            u.GetId(),
		Email:         "", //email remains the same
		Balance:       1000.0,
		ChangeBalance: false, // balance too
	}

	nu, err = grpcClient.UpdateUser(context.TODO(), updateParams)
	s.A.NotNil(err)
	s.A.Equal(codes.Canceled, grpc.Code(err))
	s.A.Nil(nu)

	// db error
	db.Engine.DropTables(&user.Users{})
	updateParams = &userGen.UpdateUserParams{
		Id:            u.GetId(),
		Email:         "user@domain.com", //email remains the same
		Balance:       1000.0,
		ChangeBalance: true, // balance too
	}

	nu, err = grpcClient.UpdateUser(context.TODO(), updateParams)
	s.A.NotNil(err)
	s.A.Equal(codes.Unknown, grpc.Code(err))
	s.A.Nil(nu)
}

func (s *ServerFunctionalSuite) TestVerifyUser() {
	db.Engine.CreateTables(&user.Users{})
	// parse error
	v := &userGen.Verify{
		Guid: "aguid",
		Id:   "fdgdfsg",
	}

	u, err := grpcClient.VerifyUser(context.TODO(), v)
	s.A.False(u.GetCanUserFund())
	s.A.Equal(codes.NotFound, grpc.Code(err))

	//success
	uss := &user.Users{ // user fixture
		Email:   "user42@example.com",
		Balance: 100.00,
	}

	db.Engine.InsertOne(uss)

	v = &userGen.Verify{
		Guid:   "another guid",
		Id:     fmt.Sprintf("%d", uss.ID),
		Amount: 50.00,
	}

	u, err = grpcClient.VerifyUser(context.TODO(), v)
	s.A.True(u.GetCanUserFund())
	s.A.Equal(codes.OK, grpc.Code(err))

	//user can't fund
	v = &userGen.Verify{
		Guid:   "another guid",
		Id:     fmt.Sprintf("%d", uss.ID),
		Amount: 101.00,
	}

	u, err = grpcClient.VerifyUser(context.TODO(), v)
	s.A.False(u.GetCanUserFund())
	s.A.Equal(codes.OK, grpc.Code(err))

	//user not found
	v = &userGen.Verify{
		Guid:   "another guid",
		Id:     "23434234", //not in test db
		Amount: 101.00,
	}

	u, err = grpcClient.VerifyUser(context.TODO(), v)
	s.A.False(u.GetCanUserFund())
	s.A.Equal(codes.NotFound, grpc.Code(err))

	//db error
	db.Engine.DropTables(&user.Users{})
	u, err = grpcClient.VerifyUser(context.TODO(), v)
	s.A.False(u.GetCanUserFund())
	s.A.Equal(codes.Unknown, grpc.Code(err))
}
<commit_msg>linter fix<commit_after>package server

import (
	"context"
	"fmt"

	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"

	log "github.com/Sirupsen/logrus"
	"github.com/agave/ah-microservices/services/users/db"
	"github.com/agave/ah-microservices/services/users/user"
	userGen "github.com/agave/ah-microservices/services/users/user/generated"
	"github.com/agave/ah-microservices/services/users/util"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/suite"
)

type ServerFunctionalSuite struct {
	suite.Suite
	A *assert.Assertions
}

func (s *ServerFunctionalSuite) SetupSuite() {
	var err error
	s.A = assert.New(s.T())

	// db
	util.Config.DB.DBName = util.Config.DB.DBName + "_test"

	url := util.FormatDBURL("postgres", &util.Config.DB)
	err = db.InitDB("postgres", url)
	if err != nil {
		log.WithFields(log.Fields{
			"error": err, "url": url,
		}).Fatalln("Error starting database")
		return
	}

	err = db.Engine.CreateTables(&user.Users{})
	if err != nil {
		log.WithField("error", err).Fatalln("Error creating tables")
	}
}

func (s *ServerFunctionalSuite) TearDownSuite() {
	if ok, _ := db.Engine.IsTableExist(&user.Users{}); ok {
		db.Engine.DropTables(&user.Users{})
	}
	if ok, _ := db.Engine.IsTableExist(&user.HeldBalance{}); ok {
		db.Engine.DropTables(&user.HeldBalance{})
	}
	db.Engine.Close()
}

func (s *ServerFunctionalSuite) TestGetUser() {
	db.Engine.CreateTables(user.Users{})
	type testCases struct {
		Ins  []*userGen.Id
		Outs []codes.Code
	}

	h := &userGen.Create{Email: "admin@admin.com", Balance: 0}
	prof, _ := grpcClient.CreateUser(context.TODO(), h)

	tc := testCases{
		Ins: []*userGen.Id{
			&userGen.Id{Guid: "1", Id: prof.GetId()}, // happy
			&userGen.Id{Guid: "2", Id: "10000"},      // not found
			&userGen.Id{Guid: "3", Id: "-99999"},
		},
		Outs: []codes.Code{
			codes.OK,
			codes.NotFound,
			codes.NotFound,
		},
	}

	for i, v := range tc.Ins {
		p, err := grpcClient.GetUser(context.TODO(), v)
		s.A.Equal(tc.Outs[i], grpc.Code(err), "GetUser should return expected error")
		if tc.Outs[i] == codes.OK {
			s.A.Equal(v.GetId(), p.GetId(), "Should return same id")
			s.A.NotEmpty(p.GetEmail(), "Email shouldn't be empty")
		}
	}

	//trigger db error
	db.Engine.DropTables(&user.Users{})
	p, err := grpcClient.GetUser(context.TODO(), &userGen.Id{Id: "1"})
	s.A.Nil(p)
	s.A.Equal(codes.Internal, grpc.Code(err))
}

func (s *ServerFunctionalSuite) TestCreateUser() {
	db.Engine.CreateTables(user.Users{})
	type testCases struct {
		Ins  []*userGen.Create
		Outs []codes.Code
	}

	tc := testCases{
		Ins: []*userGen.Create{
			&userGen.Create{Email: "admin@admin.com", Balance: 0},
			&userGen.Create{Email: "admin@admin.com", Balance: 0},
			&userGen.Create{Email: "admin2@admin.com", Balance: -42},
			&userGen.Create{Email: "admin3@admin.com", Balance: 42},
			&userGen.Create{Email: "admin4@admin.com", Balance: 3456.3453},
			&userGen.Create{Email: "", Balance: 3456.3453},
		},
		Outs: []codes.Code{
			codes.OK,
			codes.AlreadyExists,
			codes.InvalidArgument,
			codes.OK,
			codes.OK,
			codes.InvalidArgument,
		},
	}

	for i, c := range tc.Ins {
		p, err := grpcClient.CreateUser(context.TODO(), c)
		s.A.Equal(tc.Outs[i], grpc.Code(err), "Create should return expected error")
		if tc.Outs[i] == codes.OK {
			s.A.NotNil(p)
			s.A.Equal(c.GetEmail(), p.GetEmail(), "Insert should return same Email")
			s.A.Equal(c.GetBalance(), p.GetBalance(), "Insert should return same Balance")
		}
	}

	//trigger db error
	db.Engine.DropTables(&user.Users{})
	crt := &userGen.Create{
		Email:   "not empty",
		Balance: 100.00,
	}
	p, err := grpcClient.CreateUser(context.TODO(), crt)
	s.A.Nil(p)
	s.A.Equal(codes.Unknown, grpc.Code(err))
}

func (s *ServerFunctionalSuite) TestUpdateUser() {
	db.Engine.CreateTables(user.Users{})
	//insert fixtures
	createFixture := &userGen.Create{
		Email:   "user@domain.org",
		Balance: 10.00,
	}
	u, err := grpcClient.CreateUser(context.TODO(), createFixture)
	s.A.Nil(err)
	s.A.Equal(createFixture.GetEmail(), u.GetEmail(), "Should return inserted fixture's email")
	s.A.Equal(createFixture.GetBalance(), u.GetBalance(), "Should return inserted fixture's balance")
	s.A.NotEmpty(u.GetId(), "Should return with a valid Id")

	//change email
	updateParams := &userGen.UpdateUserParams{
		Id:            u.GetId(),
		Email:         "user@domain.com",
		Balance:       0.0,
		ChangeBalance: false,
	}

	nu, err := grpcClient.UpdateUser(context.TODO(), updateParams)
	s.A.Nil(err)
	s.A.Equal(updateParams.GetEmail(), nu.GetEmail(), "User email should be updated when not empty")
	s.A.NotEqual(updateParams.GetBalance(), nu.GetBalance(), "Balance shouldn't be updated when ChangeBalanceFlag is false")
	s.A.Equal(updateParams.GetId(), nu.GetId(), "Id never changes")

	// change balance
	updateParams = &userGen.UpdateUserParams{
		Id:            u.GetId(),
		Email:         "", //email remains the same
		Balance:       1000.0,
		ChangeBalance: true,
	}

	nu, err = grpcClient.UpdateUser(context.TODO(), updateParams)
	s.A.Nil(err)
	s.A.NotEqual(updateParams.GetEmail(), nu.GetEmail(), "User email should not update when empty")
	s.A.Equal(updateParams.GetBalance(), nu.GetBalance(), "Balance should be updated")

	// invalid/non existent id
	updateParams = &userGen.UpdateUserParams{
		Id:            "-666",
		Email:         "", //email remains the same
		Balance:       1000.0,
		ChangeBalance: true,
	}

	nu, err = grpcClient.UpdateUser(context.TODO(), updateParams)
	s.A.NotNil(err)
	s.A.Equal(codes.NotFound, grpc.Code(err))
	s.A.Nil(nu)

	//no changes
	updateParams = &userGen.UpdateUserParams{
		Id:            u.GetId(),
		Email:         "", //email remains the same
		Balance:       1000.0,
		ChangeBalance: false, // balance too
	}

	nu, err = grpcClient.UpdateUser(context.TODO(), updateParams)
	s.A.NotNil(err)
	s.A.Equal(codes.Canceled, grpc.Code(err))
	s.A.Nil(nu)

	// db error
	db.Engine.DropTables(&user.Users{})
	updateParams = &userGen.UpdateUserParams{
		Id:            u.GetId(),
		Email:         "user@domain.com", //email remains the same
		Balance:       1000.0,
		ChangeBalance: true, // balance too
	}

	nu, err = grpcClient.UpdateUser(context.TODO(), updateParams)
	s.A.NotNil(err)
	s.A.Equal(codes.Unknown, grpc.Code(err))
	s.A.Nil(nu)
}

func (s *ServerFunctionalSuite) TestVerifyUser() {
	db.Engine.CreateTables(&user.Users{})
	// parse error
	v := &userGen.Verify{
		Guid: "aguid",
		Id:   "fdgdfsg",
	}

	u, err := grpcClient.VerifyUser(context.TODO(), v)
	s.A.False(u.GetCanUserFund())
	s.A.Equal(codes.NotFound, grpc.Code(err))

	//success
	uss := &user.Users{ // user fixture
		Email:   "user42@example.com",
		Balance: 100.00,
	}

	db.Engine.InsertOne(uss)

	v = &userGen.Verify{
		Guid:   "another guid",
		Id:     fmt.Sprintf("%d", uss.ID),
		Amount: 50.00,
	}

	u, err = grpcClient.VerifyUser(context.TODO(), v)
	s.A.True(u.GetCanUserFund())
	s.A.Equal(codes.OK, grpc.Code(err))

	//user can't fund
	v = &userGen.Verify{
		Guid:   "another guid",
		Id:     fmt.Sprintf("%d", uss.ID),
		Amount: 101.00,
	}

	u, err = grpcClient.VerifyUser(context.TODO(), v)
	s.A.False(u.GetCanUserFund())
	s.A.Equal(codes.OK, grpc.Code(err))

	//user not found
	v = &userGen.Verify{
		Guid:   "another guid",
		Id:     "23434234", //not in test db
		Amount: 101.00,
	}

	u, err = grpcClient.VerifyUser(context.TODO(), v)
	s.A.False(u.GetCanUserFund())
	s.A.Equal(codes.NotFound, grpc.Code(err))

	//db error
	db.Engine.DropTables(&user.Users{})
	u, err = grpcClient.VerifyUser(context.TODO(), v)
	s.A.False(u.GetCanUserFund())
	s.A.Equal(codes.Unknown, grpc.Code(err))
}
<|endoftext|>


<commit_before>package ordset

import (
    "github.com/odeke-em/goht/dnode"
)

type OrdSet struct {
    dl *dnode.DNode
    jar map[uint] *dnode.DNode
}

func New() *OrdSet {
    od := new(OrdSet)
    od.dl = dnode.New(nil)
    od.jar = map[uint] *dnode.DNode {}

    return od
}

func (o *OrdSet) getDNode(h uint) *dnode.DNode {
    if o.jar == nil {
        return nil
    } else {
        retr, _ := o.jar[h]
        return retr
    }
}

func (o *OrdSet) Get(h uint) (interface{}, bool) {
    if o.jar == nil {
        return nil, false
    } else {
        retr, check := o.jar[h]
        if check == false || retr == nil{
            return nil, false
        } else {
            return retr.Value, true
        }
    }
}

func (o *OrdSet) putIntoJar(k uint, v *dnode.DNode) {
    o.jar[k] = v
}

func (o *OrdSet) Add(h uint, item interface{}) (vacated interface{}, fIns bool) {
    vacated = nil
    fIns = true // First time insert

    retr := o.getDNode(h)
    if retr == nil { // Time to add it
        nd := o.dl.Prepend(item)
        o.putIntoJar(h, nd)
    } else { // Update that value
        vacated = retr.SetValue(item)
        fIns = false
    }

    return
}

func (o *OrdSet) Remove(h uint) (vacated interface{}, exists bool) {
    retr := o.getDNode(h)

    if retr == nil {
        return nil, false
    } else {
        prev := retr.GetPrev()
        next := retr.GetNext()

        if prev != nil {
            prev.SetNext(next)
        }

        if next != nil {
            next.SetPrev(prev)
        }

        delete(o.jar, h)
        return retr.Value, true
    }
}

func (o *OrdSet) Iterator() *dnode.DNode {
    return o.dl
}
<commit_msg>Added description<commit_after>/*
    Author: Emmanuel Odeke <odeke@ualberta.ca>
    Description: Data structure that allows for set operations:
        put, pop, get
    however in O(1) time.
    The caveat however is that extra memory is used since under the
    hood, it consists of a doubly linked list as well as a hashmap.
    The hashmap is a item-hash to doubly linked list node map.
*/

package ordset

import (
    "github.com/odeke-em/goht/dnode"
)

type OrdSet struct {
    dl *dnode.DNode
    jar map[uint] *dnode.DNode
}

func New() *OrdSet {
    od := new(OrdSet)
    od.dl = dnode.New(nil)
    od.jar = map[uint] *dnode.DNode {}

    return od
}

func (o *OrdSet) getDNode(h uint) *dnode.DNode {
    if o.jar == nil {
        return nil
    } else {
        retr, _ := o.jar[h]
        return retr
    }
}

func (o *OrdSet) Get(h uint) (interface{}, bool) {
    if o.jar == nil {
        return nil, false
    } else {
        retr, check := o.jar[h]
        if check == false || retr == nil{
            return nil, false
        } else {
            return retr.Value, true
        }
    }
}

func (o *OrdSet) putIntoJar(k uint, v *dnode.DNode) {
    o.jar[k] = v
}

func (o *OrdSet) Add(h uint, item interface{}) (vacated interface{}, fIns bool) {
    vacated = nil
    fIns = true // First time insert

    retr := o.getDNode(h)
    if retr == nil { // Time to add it
        nd := o.dl.Prepend(item)
        o.putIntoJar(h, nd)
    } else { // Update that value
        vacated = retr.SetValue(item)
        fIns = false
    }

    return
}

func (o *OrdSet) Remove(h uint) (vacated interface{}, exists bool) {
    retr := o.getDNode(h)

    if retr == nil {
        return nil, false
    } else {
        prev := retr.GetPrev()
        next := retr.GetNext()

        if prev != nil {
            prev.SetNext(next)
        }

        if next != nil {
            next.SetPrev(prev)
        }

        delete(o.jar, h)
        return retr.Value, true
    }
}

func (o *OrdSet) Iterator() *dnode.DNode {
    return o.dl
}
<|endoftext|>


<commit_before>package main

func PrimeFactors(number uint64) []uint64 {
	results := make([]uint64, 0, 16)

	primes := Primes(number)

	for _, prime := range primes {
		if number%prime == 0 {
			results = append(results, prime)
		}
	}

	if len(results) == 0 {
		return append(results, number)
	}

	return results
}

//Primes returns all prime factor candidates for n. i.e all primes x such that x * x < ns
func Primes(n uint64) []uint64 {
	if n < 2 {
		return nil
	}
	primes := make([]uint64, 1, 1024)
	primes[0] = 2

	for x := uint64(3); x*x <= n; x += 2 {
		prime := true
		for _, v := range primes {
			if x%v == 0 {
				prime = false
				break
			}
		}

		if prime {
			primes = append(primes, x)
		}
	}

	return primes
}
<commit_msg>2000x time speed up on euler for by fixing algorithmic complexity of prime candidacy check<commit_after>package main

func PrimeFactors(number uint64) []uint64 {
	results := make([]uint64, 0, 16)

	primes := Primes(number)

	for _, prime := range primes {
		if number%prime == 0 {
			results = append(results, prime)
		}
	}

	if len(results) == 0 {
		return append(results, number)
	}

	return results
}

//Primes returns all prime factor candidates for n. i.e all primes x such that x * x < ns
func Primes(n uint64) []uint64 {
	if n < 2 {
		return nil
	}
	primes := make([]uint64, 1, 1024)
	primes[0] = 2

	for x := uint64(3); x*x <= n; x += 2 {
		prime := true
		for i := 1; i < len(primes) && primes[i]*primes[i] <= x; i++ {
			if x%primes[i] == 0 {
				prime = false
				break
			}
		}

		if prime {
			primes = append(primes, x)
		}
	}

	return primes
}
<|endoftext|>


<commit_before>// errchk $G $D/$F.go

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

package main

const none = 0  // same const identifier declared twice should not be accepted
const none = 1  // ERROR "redeclared"
<commit_msg>Tweak comments so that this test passes with gccgo.<commit_after>// errchk $G $D/$F.go

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

package main

// same const identifier declared twice should not be accepted
const none = 0  // GCCGO_ERROR "previous"
const none = 1  // ERROR "redeclared|redef"
<|endoftext|>


<commit_before>package nsq_consumer

import (
	"context"
	"sync"

	"github.com/influxdata/telegraf"
	"github.com/influxdata/telegraf/plugins/inputs"
	"github.com/influxdata/telegraf/plugins/parsers"
	nsq "github.com/nsqio/go-nsq"
)

const (
	defaultMaxUndeliveredMessages = 1000
)

type empty struct{}
type semaphore chan empty

type logger struct {
	log telegraf.Logger
}

func (l *logger) Output(_ int, s string) error {
	l.log.Debug(s)
	return nil
}

//NSQConsumer represents the configuration of the plugin
type NSQConsumer struct {
	Server      string   `toml:"server"`
	Nsqd        []string `toml:"nsqd"`
	Nsqlookupd  []string `toml:"nsqlookupd"`
	Topic       string   `toml:"topic"`
	Channel     string   `toml:"channel"`
	MaxInFlight int      `toml:"max_in_flight"`

	MaxUndeliveredMessages int `toml:"max_undelivered_messages"`

	parser   parsers.Parser
	consumer *nsq.Consumer

	Log telegraf.Logger

	mu       sync.Mutex
	messages map[telegraf.TrackingID]*nsq.Message
	wg       sync.WaitGroup
	cancel   context.CancelFunc
}

var sampleConfig = `
  ## Server option still works but is deprecated, we just prepend it to the nsqd array.
  # server = "localhost:4150"

  ## An array representing the NSQD TCP HTTP Endpoints
  nsqd = ["localhost:4150"]

  ## An array representing the NSQLookupd HTTP Endpoints
  nsqlookupd = ["localhost:4161"]
  topic = "telegraf"
  channel = "consumer"
  max_in_flight = 100

  ## Maximum messages to read from the broker that have not been written by an
  ## output.  For best throughput set based on the number of metrics within
  ## each message and the size of the output's metric_batch_size.
  ##
  ## For example, if each message from the queue contains 10 metrics and the
  ## output metric_batch_size is 1000, setting this to 100 will ensure that a
  ## full batch is collected and the write is triggered immediately without
  ## waiting until the next flush_interval.
  # max_undelivered_messages = 1000

  ## Data format to consume.
  ## Each data format has its own unique set of configuration options, read
  ## more about them here:
  ## https://github.com/influxdata/telegraf/blob/master/docs/DATA_FORMATS_INPUT.md
  data_format = "influx"
`

// SetParser takes the data_format from the config and finds the right parser for that format
func (n *NSQConsumer) SetParser(parser parsers.Parser) {
	n.parser = parser
}

// SampleConfig returns config values for generating a sample configuration file
func (n *NSQConsumer) SampleConfig() string {
	return sampleConfig
}

// Description prints description string
func (n *NSQConsumer) Description() string {
	return "Read NSQ topic for metrics."
}

// Start pulls data from nsq
func (n *NSQConsumer) Start(ac telegraf.Accumulator) error {
	acc := ac.WithTracking(n.MaxUndeliveredMessages)
	sem := make(semaphore, n.MaxUndeliveredMessages)
	n.messages = make(map[telegraf.TrackingID]*nsq.Message, n.MaxUndeliveredMessages)

	ctx, cancel := context.WithCancel(context.Background())
	n.cancel = cancel

	if err := n.connect(); err != nil {
		return err
	}
	n.consumer.SetLogger(&logger{log: n.Log}, nsq.LogLevelInfo)
	n.consumer.AddHandler(nsq.HandlerFunc(func(message *nsq.Message) error {
		metrics, err := n.parser.Parse(message.Body)
		if err != nil {
			acc.AddError(err)
			// Remove the message from the queue
			message.Finish()
			return nil
		}
		if len(metrics) == 0 {
			message.Finish()
			return nil
		}

		select {
		case <-ctx.Done():
			return ctx.Err()
		case sem <- empty{}:
			break
		}

		n.mu.Lock()
		id := acc.AddTrackingMetricGroup(metrics)
		n.messages[id] = message
		n.mu.Unlock()
		message.DisableAutoResponse()
		return nil
	}))

	if len(n.Nsqlookupd) > 0 {
		err := n.consumer.ConnectToNSQLookupds(n.Nsqlookupd)
		if err != nil && err != nsq.ErrAlreadyConnected {
			return err
		}
	}
	err := n.consumer.ConnectToNSQDs(append(n.Nsqd, n.Server))
	if err != nil && err != nsq.ErrAlreadyConnected {
		return err
	}

	n.wg.Add(1)
	go func() {
		defer n.wg.Done()
		n.onDelivery(ctx, acc, sem)
	}()
	return nil
}

func (n *NSQConsumer) onDelivery(ctx context.Context, acc telegraf.TrackingAccumulator, sem semaphore) {
	for {
		select {
		case <-ctx.Done():
			return
		case info := <-acc.Delivered():
			n.mu.Lock()
			msg, ok := n.messages[info.ID()]
			if !ok {
				n.mu.Unlock()
				continue
			}
			<-sem
			delete(n.messages, info.ID())
			n.mu.Unlock()

			if info.Delivered() {
				msg.Finish()
			} else {
				msg.Requeue(-1)
			}
		}
	}
}

// Stop processing messages
func (n *NSQConsumer) Stop() {
	n.cancel()
	n.wg.Wait()
	n.consumer.Stop()
	<-n.consumer.StopChan
}

// Gather is a noop
func (n *NSQConsumer) Gather(_ telegraf.Accumulator) error {
	return nil
}

func (n *NSQConsumer) connect() error {
	if n.consumer == nil {
		config := nsq.NewConfig()
		config.MaxInFlight = n.MaxInFlight
		consumer, err := nsq.NewConsumer(n.Topic, n.Channel, config)
		if err != nil {
			return err
		}
		n.consumer = consumer
	}
	return nil
}

func init() {
	inputs.Add("nsq_consumer", func() telegraf.Input {
		return &NSQConsumer{
			MaxUndeliveredMessages: defaultMaxUndeliveredMessages,
		}
	})
}
<commit_msg>Fix attempt to connect to an empty list of servers. (#9503)<commit_after>package nsq_consumer

import (
	"context"
	"fmt"
	"sync"

	"github.com/influxdata/telegraf"
	"github.com/influxdata/telegraf/plugins/inputs"
	"github.com/influxdata/telegraf/plugins/parsers"
	nsq "github.com/nsqio/go-nsq"
)

const (
	defaultMaxUndeliveredMessages = 1000
)

type empty struct{}
type semaphore chan empty

type logger struct {
	log telegraf.Logger
}

func (l *logger) Output(_ int, s string) error {
	l.log.Debug(s)
	return nil
}

//NSQConsumer represents the configuration of the plugin
type NSQConsumer struct {
	Server      string   `toml:"server"`
	Nsqd        []string `toml:"nsqd"`
	Nsqlookupd  []string `toml:"nsqlookupd"`
	Topic       string   `toml:"topic"`
	Channel     string   `toml:"channel"`
	MaxInFlight int      `toml:"max_in_flight"`

	MaxUndeliveredMessages int `toml:"max_undelivered_messages"`

	parser   parsers.Parser
	consumer *nsq.Consumer

	Log telegraf.Logger

	mu       sync.Mutex
	messages map[telegraf.TrackingID]*nsq.Message
	wg       sync.WaitGroup
	cancel   context.CancelFunc
}

var sampleConfig = `
  ## Server option still works but is deprecated, we just prepend it to the nsqd array.
  # server = "localhost:4150"

  ## An array representing the NSQD TCP HTTP Endpoints
  nsqd = ["localhost:4150"]

  ## An array representing the NSQLookupd HTTP Endpoints
  nsqlookupd = ["localhost:4161"]
  topic = "telegraf"
  channel = "consumer"
  max_in_flight = 100

  ## Maximum messages to read from the broker that have not been written by an
  ## output.  For best throughput set based on the number of metrics within
  ## each message and the size of the output's metric_batch_size.
  ##
  ## For example, if each message from the queue contains 10 metrics and the
  ## output metric_batch_size is 1000, setting this to 100 will ensure that a
  ## full batch is collected and the write is triggered immediately without
  ## waiting until the next flush_interval.
  # max_undelivered_messages = 1000

  ## Data format to consume.
  ## Each data format has its own unique set of configuration options, read
  ## more about them here:
  ## https://github.com/influxdata/telegraf/blob/master/docs/DATA_FORMATS_INPUT.md
  data_format = "influx"
`

// SetParser takes the data_format from the config and finds the right parser for that format
func (n *NSQConsumer) SetParser(parser parsers.Parser) {
	n.parser = parser
}

// SampleConfig returns config values for generating a sample configuration file
func (n *NSQConsumer) SampleConfig() string {
	return sampleConfig
}

// Description prints description string
func (n *NSQConsumer) Description() string {
	return "Read NSQ topic for metrics."
}

// Start pulls data from nsq
func (n *NSQConsumer) Start(ac telegraf.Accumulator) error {
	acc := ac.WithTracking(n.MaxUndeliveredMessages)
	sem := make(semaphore, n.MaxUndeliveredMessages)
	n.messages = make(map[telegraf.TrackingID]*nsq.Message, n.MaxUndeliveredMessages)

	ctx, cancel := context.WithCancel(context.Background())
	n.cancel = cancel

	if err := n.connect(); err != nil {
		return err
	}
	n.consumer.SetLogger(&logger{log: n.Log}, nsq.LogLevelInfo)
	n.consumer.AddHandler(nsq.HandlerFunc(func(message *nsq.Message) error {
		metrics, err := n.parser.Parse(message.Body)
		if err != nil {
			acc.AddError(err)
			// Remove the message from the queue
			message.Finish()
			return nil
		}
		if len(metrics) == 0 {
			message.Finish()
			return nil
		}

		select {
		case <-ctx.Done():
			return ctx.Err()
		case sem <- empty{}:
			break
		}

		n.mu.Lock()
		id := acc.AddTrackingMetricGroup(metrics)
		n.messages[id] = message
		n.mu.Unlock()
		message.DisableAutoResponse()
		return nil
	}))

	// For backward compatibility
	if n.Server != "" {
		n.Nsqd = append(n.Nsqd, n.Server)
	}

	// Check if we have anything to connect to
	if len(n.Nsqlookupd) == 0 && len(n.Nsqd) == 0 {
		return fmt.Errorf("either 'nsqd' or 'nsqlookupd' needs to be specified")
	}

	if len(n.Nsqlookupd) > 0 {
		err := n.consumer.ConnectToNSQLookupds(n.Nsqlookupd)
		if err != nil && err != nsq.ErrAlreadyConnected {
			return err
		}
	}

	if len(n.Nsqd) > 0 {
		err := n.consumer.ConnectToNSQDs(n.Nsqd)
		if err != nil && err != nsq.ErrAlreadyConnected {
			return err
		}
	}

	n.wg.Add(1)
	go func() {
		defer n.wg.Done()
		n.onDelivery(ctx, acc, sem)
	}()
	return nil
}

func (n *NSQConsumer) onDelivery(ctx context.Context, acc telegraf.TrackingAccumulator, sem semaphore) {
	for {
		select {
		case <-ctx.Done():
			return
		case info := <-acc.Delivered():
			n.mu.Lock()
			msg, ok := n.messages[info.ID()]
			if !ok {
				n.mu.Unlock()
				continue
			}
			<-sem
			delete(n.messages, info.ID())
			n.mu.Unlock()

			if info.Delivered() {
				msg.Finish()
			} else {
				msg.Requeue(-1)
			}
		}
	}
}

// Stop processing messages
func (n *NSQConsumer) Stop() {
	n.cancel()
	n.wg.Wait()
	n.consumer.Stop()
	<-n.consumer.StopChan
}

// Gather is a noop
func (n *NSQConsumer) Gather(_ telegraf.Accumulator) error {
	return nil
}

func (n *NSQConsumer) connect() error {
	if n.consumer == nil {
		config := nsq.NewConfig()
		config.MaxInFlight = n.MaxInFlight
		consumer, err := nsq.NewConsumer(n.Topic, n.Channel, config)
		if err != nil {
			return err
		}
		n.consumer = consumer
	}
	return nil
}

func init() {
	inputs.Add("nsq_consumer", func() telegraf.Input {
		return &NSQConsumer{
			MaxUndeliveredMessages: defaultMaxUndeliveredMessages,
		}
	})
}
<|endoftext|>


<commit_before>package common

import (
	"context"
	"crypto/tls"
	"fmt"
	"net/http"

	"github.com/hashicorp/packer/helper/multistep"
	"github.com/hashicorp/packer/packer"
	"github.com/hashicorp/packer/template/interpolate"
	"github.com/outscale/osc-go/oapi"
)

type StepUpdateOMIAttributes struct {
	AccountIds         []string
	SnapshotAccountIds []string
	Ctx                interpolate.Context
}

func (s *StepUpdateOMIAttributes) Run(_ context.Context, state multistep.StateBag) multistep.StepAction {
	oapiconn := state.Get("oapi").(*oapi.Client)
	config := state.Get("clientConfig").(*oapi.Config)
	ui := state.Get("ui").(packer.Ui)
	omis := state.Get("omis").(map[string]string)
	snapshots := state.Get("snapshots").(map[string][]string)

	// Determine if there is any work to do.
	valid := false
	valid = valid || (s.AccountIds != nil && len(s.AccountIds) > 0)
	valid = valid || (s.SnapshotAccountIds != nil && len(s.SnapshotAccountIds) > 0)

	if !valid {
		return multistep.ActionContinue
	}

	s.Ctx.Data = extractBuildInfo(oapiconn.GetConfig().Region, state)

	updateSnapshoptRequest := oapi.UpdateSnapshotRequest{
		PermissionsToCreateVolume: oapi.PermissionsOnResourceCreation{
			Additions: oapi.PermissionsOnResource{
				AccountIds:       s.AccountIds,
				GlobalPermission: true,
			},
		},
	}

	updateImageRequest := oapi.UpdateImageRequest{
		PermissionsToLaunch: oapi.PermissionsOnResourceCreation{
			Additions: oapi.PermissionsOnResource{
				AccountIds:       s.AccountIds,
				GlobalPermission: true,
			},
		},
	}

	// Updating image attributes
	for region, omi := range omis {
		ui.Say(fmt.Sprintf("Updating attributes on OMI (%s)...", omi))
		newConfig := &oapi.Config{
			UserAgent: config.UserAgent,
			SecretKey: config.SecretKey,
			Service:   config.Service,
			Region:    region, //New region
			URL:       config.URL,
		}

		skipClient := &http.Client{
			Transport: &http.Transport{
				TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
			},
		}

		regionconn := oapi.NewClient(newConfig, skipClient)

		ui.Message(fmt.Sprintf("Updating: %s", omi))
		updateImageRequest.ImageId = omi
		_, err := regionconn.POST_UpdateImage(updateImageRequest)
		if err != nil {
			err := fmt.Errorf("Error updating OMI: %s", err)
			state.Put("error", err)
			ui.Error(err.Error())
			return multistep.ActionHalt
		}
	}

	// Updating snapshot attributes
	for region, region_snapshots := range snapshots {
		for _, snapshot := range region_snapshots {
			ui.Say(fmt.Sprintf("Updating attributes on snapshot (%s)...", snapshot))
			newConfig := &oapi.Config{
				UserAgent: config.UserAgent,
				SecretKey: config.SecretKey,
				Service:   config.Service,
				Region:    region, //New region
				URL:       config.URL,
			}

			skipClient := &http.Client{
				Transport: &http.Transport{
					TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
				},
			}

			regionconn := oapi.NewClient(newConfig, skipClient)

			ui.Message(fmt.Sprintf("Updating: %s", snapshot))
			updateSnapshoptRequest.SnapshotId = snapshot
			_, err := regionconn.POST_UpdateSnapshot(updateSnapshoptRequest)
			if err != nil {
				err := fmt.Errorf("Error updating snapshot: %s", err)
				state.Put("error", err)
				ui.Error(err.Error())
				return multistep.ActionHalt
			}

		}
	}

	return multistep.ActionContinue
}

func (s *StepUpdateOMIAttributes) Cleanup(state multistep.StateBag) {
	// No cleanup...
}
<commit_msg>fix: add missing secret key in http client configuration<commit_after>package common

import (
	"context"
	"crypto/tls"
	"fmt"
	"net/http"

	"github.com/hashicorp/packer/helper/multistep"
	"github.com/hashicorp/packer/packer"
	"github.com/hashicorp/packer/template/interpolate"
	"github.com/outscale/osc-go/oapi"
)

type StepUpdateOMIAttributes struct {
	AccountIds         []string
	SnapshotAccountIds []string
	Ctx                interpolate.Context
}

func (s *StepUpdateOMIAttributes) Run(_ context.Context, state multistep.StateBag) multistep.StepAction {
	oapiconn := state.Get("oapi").(*oapi.Client)
	config := state.Get("clientConfig").(*oapi.Config)
	ui := state.Get("ui").(packer.Ui)
	omis := state.Get("omis").(map[string]string)
	snapshots := state.Get("snapshots").(map[string][]string)

	// Determine if there is any work to do.
	valid := false
	valid = valid || (s.AccountIds != nil && len(s.AccountIds) > 0)
	valid = valid || (s.SnapshotAccountIds != nil && len(s.SnapshotAccountIds) > 0)

	if !valid {
		return multistep.ActionContinue
	}

	s.Ctx.Data = extractBuildInfo(oapiconn.GetConfig().Region, state)

	updateSnapshoptRequest := oapi.UpdateSnapshotRequest{
		PermissionsToCreateVolume: oapi.PermissionsOnResourceCreation{
			Additions: oapi.PermissionsOnResource{
				AccountIds:       s.AccountIds,
				GlobalPermission: true,
			},
		},
	}

	updateImageRequest := oapi.UpdateImageRequest{
		PermissionsToLaunch: oapi.PermissionsOnResourceCreation{
			Additions: oapi.PermissionsOnResource{
				AccountIds:       s.AccountIds,
				GlobalPermission: true,
			},
		},
	}

	// Updating image attributes
	for region, omi := range omis {
		ui.Say(fmt.Sprintf("Updating attributes on OMI (%s)...", omi))
		newConfig := &oapi.Config{
			UserAgent: config.UserAgent,
			AccessKey: config.AccessKey,
			SecretKey: config.SecretKey,
			Service:   config.Service,
			Region:    region, //New region
			URL:       config.URL,
		}

		skipClient := &http.Client{
			Transport: &http.Transport{
				TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
			},
		}

		regionconn := oapi.NewClient(newConfig, skipClient)

		ui.Message(fmt.Sprintf("Updating: %s", omi))
		updateImageRequest.ImageId = omi
		_, err := regionconn.POST_UpdateImage(updateImageRequest)
		if err != nil {
			err := fmt.Errorf("Error updating OMI: %s", err)
			state.Put("error", err)
			ui.Error(err.Error())
			return multistep.ActionHalt
		}
	}

	// Updating snapshot attributes
	for region, region_snapshots := range snapshots {
		for _, snapshot := range region_snapshots {
			ui.Say(fmt.Sprintf("Updating attributes on snapshot (%s)...", snapshot))
			newConfig := &oapi.Config{
				UserAgent: config.UserAgent,
				SecretKey: config.SecretKey,
				Service:   config.Service,
				Region:    region, //New region
				URL:       config.URL,
			}

			skipClient := &http.Client{
				Transport: &http.Transport{
					TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
				},
			}

			regionconn := oapi.NewClient(newConfig, skipClient)

			ui.Message(fmt.Sprintf("Updating: %s", snapshot))
			updateSnapshoptRequest.SnapshotId = snapshot
			_, err := regionconn.POST_UpdateSnapshot(updateSnapshoptRequest)
			if err != nil {
				err := fmt.Errorf("Error updating snapshot: %s", err)
				state.Put("error", err)
				ui.Error(err.Error())
				return multistep.ActionHalt
			}

		}
	}

	return multistep.ActionContinue
}

func (s *StepUpdateOMIAttributes) Cleanup(state multistep.StateBag) {
	// No cleanup...
}
<|endoftext|>


<commit_before>package main

import (
	"io"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/Sirupsen/logrus"
	"github.com/juju/ratelimit"
	"github.com/pkg/errors"
	"github.com/tomnomnom/linkheader"
)

func PubSubMakeCallbackURL(baseURL, id string) string {
	return baseURL + "/" + id
}

type PubSubSubscription struct {
	ID          string
	Hub         string
	Topic       string
	CallbackURL string
	CreatedAt   time.Time
	UpdatedAt   time.Time
	ExpiresAt   *time.Time
}

type PubSubState interface {
	All() (subscriptions []PubSubSubscription, err error)
	Add(hub, topic, baseURL string) (subscription *PubSubSubscription, oldCallbackURL string, err error)
	Get(hub, topic string) (subscription *PubSubSubscription, err error)
	GetByID(id string) (subscription *PubSubSubscription, err error)
	Set(hub, topic string, updatedAt, expiresAt time.Time) (err error)
	Del(hub, topic string) (err error)
}

type PubSubMessageHandler func(id string, s *PubSubSubscription, rd io.ReadCloser)

type PubSubClient struct {
	CallbackURL string
	State       PubSubState
	OnMessage   PubSubMessageHandler
}

func NewPubSubClient(callbackURL string, state PubSubState, onMessage PubSubMessageHandler) *PubSubClient {
	return &PubSubClient{
		CallbackURL: callbackURL,
		State:       state,
		OnMessage:   onMessage,
	}
}

func (c *PubSubClient) Refresh(forceUpdate bool, interval time.Duration) error {
	a, err := c.State.All()
	if err != nil {
		return errors.Wrap(err, "PubSubClient.Refresh")
	}

	logrus.WithField("count", len(a)).Debug("pubsub: got subscriptions to refresh")

	m := make(map[string]*ratelimit.Bucket)
	var l sync.RWMutex

	getBucket := func(host string) *ratelimit.Bucket {
		l.RLock()
		if b, ok := m[host]; ok {
			l.RUnlock()
			return b
		}
		l.RUnlock()

		l.Lock()
		if b, ok := m[host]; ok {
			l.Unlock()
			return b
		}
		defer l.Unlock()

		m[host] = ratelimit.NewBucket(time.Second*30, 40000)

		return m[host]
	}

	var g WorkerGroup

	for _, e := range a {
		e := e

		callbackURL := c.CallbackURL + "/" + e.ID

		l := logrus.WithFields(logrus.Fields{
			"id":               e.ID,
			"hub":              e.Hub,
			"topic":            e.Topic,
			"callback_url":     e.CallbackURL,
			"new_callback_url": callbackURL,
			"created_at":       e.CreatedAt,
			"updated_at":       e.UpdatedAt,
			"expires_at":       e.ExpiresAt,
			"force_update":     forceUpdate,
		})

		if e.ExpiresAt == nil {
			l.Debug("pubsub: not refreshing subscription which has never been verified")
			continue
		}

		if forceUpdate || e.ExpiresAt.Sub(time.Now()) < interval || e.CallbackURL != callbackURL {
			g.Add(func() error {
				u, err := url.Parse(e.Hub)
				if err != nil {
					l.WithError(err).Warn("pubsub: couldn't parse hub url")
					return errors.Wrap(err, "PubSubClient.RefreshWorker")
				}

				if dur, ok := getBucket(u.Host).TakeMaxDuration(10000, *pubsubRefreshInterval); !ok {
					l.Debug("pubsub: skipping renewing for now as we'd have to wait too long")
					return nil
				} else if dur > 0 {
					l.WithField("duration", dur).Debug("pubsub: waiting so as not to overwhelm the endpoint")
					time.Sleep(dur)
				}

				l.Debug("pubsub: refreshing subscription")

				// simulate subscription
				time.Sleep(time.Second * 5)

				// if err := c.Subscribe(e.Hub, e.Topic); err != nil {
				// 	l.WithError(err).Warn("pubsub: couldn't subscribe to topic")
				// 	return errors.Wrap(err, "PubSubClient.RefreshWorker")
				// }

				l.Debug("pubsub: subscribed successfully")
				return nil
			})
		}
	}

	return errors.Wrap(g.Run(4), "PubSubClient.Refresh")
}

func (c *PubSubClient) Subscribe(hub, topic string) error {
	logrus.WithFields(logrus.Fields{"hub": hub, "topic": topic}).Debug("pubsub: subscribing")

	s, oldCallbackURL, err := c.State.Add(hub, topic, c.CallbackURL)
	if err != nil {
		return errors.Wrap(err, "PubSubClient.Subscribe")
	}

	if oldCallbackURL != s.CallbackURL {
		if oldCallbackURL != "" {
			if err := PubSubUnsubscribe(hub, topic, oldCallbackURL); err != nil {
				return errors.Wrap(err, "PubSubClient.Subscribe")
			}
		}

		if err := PubSubSubscribe(hub, topic, s.CallbackURL); err != nil {
			return errors.Wrap(err, "PubSubClient.Subscribe")
		}
	}

	return nil
}

func (c *PubSubClient) Unsubscribe(hub, topic string) error {
	logrus.WithFields(logrus.Fields{"hub": hub, "topic": topic}).Debug("pubsub: unsubscribing")

	s, err := c.State.Get(hub, topic)
	if err != nil {
		return errors.Wrap(err, "PubSubClient.Unsubscribe")
	}

	if s != nil {
		if err := PubSubUnsubscribe(hub, topic, s.CallbackURL); err != nil {
			return errors.Wrap(err, "PubSubClient.Unsubscribe")
		}

		if err := c.State.Del(s.Hub, s.Topic); err != nil {
			return errors.Wrap(err, "PubSubClient.Unsubscribe")
		}
	}

	return nil
}

func (c *PubSubClient) Handler() *PubSubHandler {
	return &PubSubHandler{
		OnChallenge: func(id, topic, mode string, leaseTime time.Duration) error {
			l := logrus.WithFields(logrus.Fields{
				"id":         id,
				"topic":      topic,
				"mode":       mode,
				"lease_time": leaseTime,
			})

			l.Debug("pubsub: received challenge")

			s, err := c.State.GetByID(id)
			if err != nil {
				l.WithError(err).Warn("pubsub: error fetching subscription during challenge")
				return err
			}

			if s == nil {
				l.WithError(err).Warn("pubsub: subscription not found during challenge")
				return errors.Errorf("PubSubClient.Handler: subscription not found")
			}

			return errors.Wrap(c.State.Set(s.Hub, s.Topic, time.Now(), time.Now().Add(leaseTime)), "PubSubClient.Handler")
		},
		OnMessage: func(id, topic string, rd io.ReadCloser) {
			l := logrus.WithFields(logrus.Fields{
				"id":    id,
				"topic": topic,
			})

			l.Debug("pubsub: received message")

			defer rd.Close()

			if c.OnMessage != nil {
				s, err := c.State.GetByID(id)
				if err != nil {
					l.WithError(err).Warn("pubsub: error fetching subscription during reception")
					return
				}

				c.OnMessage(id, s, rd)
			}
		},
	}
}

func PubSubAlter(hub, topic, callbackURL, mode string) error {
	res, err := http.PostForm(hub, url.Values{
		"hub.callback":      []string{callbackURL},
		"hub.mode":          []string{mode},
		"hub.topic":         []string{topic},
		"hub.verify":        []string{"async"},
		"hub.lease_seconds": []string{"604800"},
	})

	if err != nil {
		return errors.Wrap(err, "PubSubAlter")
	}

	if res.StatusCode < 200 || res.StatusCode >= 300 {
		return errors.Errorf("PubSubAlter: invalid status code; expected 2xx but got %d", res.StatusCode)
	}

	return nil
}

func PubSubSubscribe(hub, topic, callbackURL string) error {
	return errors.Wrap(PubSubAlter(hub, topic, callbackURL, "subscribe"), "PubSubSubscribe")
}

func PubSubUnsubscribe(hub, topic, callbackURL string) error {
	return errors.Wrap(PubSubAlter(hub, topic, callbackURL, "unsubscribe"), "PubSubSubscribe")
}

type PubSubHandler struct {
	OnChallenge func(id, topic, mode string, leaseTime time.Duration) error
	OnMessage   func(id, topic string, rd io.ReadCloser)
}

func (h *PubSubHandler) ServeHTTP(rw http.ResponseWriter, r *http.Request) {
	q := r.URL.Query()

	bits := strings.Split(r.URL.Path, "/")
	id := bits[len(bits)-1]

	if r.Method == http.MethodGet && q.Get("hub.challenge") != "" {
		if h.OnChallenge != nil {
			leaseTime := time.Hour * 12
			if n, err := strconv.ParseInt(q.Get("hub.lease_seconds"), 10, 64); err == nil {
				leaseTime = time.Second * time.Duration(n)
			}

			if err := h.OnChallenge(id, q.Get("hub.topic"), q.Get("hub.mode"), leaseTime); err != nil {
				rw.WriteHeader(http.StatusNotFound)
				return
			}
		}

		rw.Write([]byte(r.URL.Query().Get("hub.challenge")))
		return
	}

	if r.Method != http.MethodPost {
		rw.WriteHeader(http.StatusMethodNotAllowed)
		return
	}

	links := linkheader.ParseMultiple(r.Header["Link"])

	var topic string
	if a := links.FilterByRel("self"); len(a) > 0 {
		topic = a[0].URL
	}

	rw.WriteHeader(http.StatusAccepted)

	if h.OnMessage != nil {
		h.OnMessage(id, topic, r.Body)
	}
}
<commit_msg>calculate wait time outside worker loop<commit_after>package main

import (
	"io"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/Sirupsen/logrus"
	"github.com/juju/ratelimit"
	"github.com/pkg/errors"
	"github.com/tomnomnom/linkheader"
)

func PubSubMakeCallbackURL(baseURL, id string) string {
	return baseURL + "/" + id
}

type PubSubSubscription struct {
	ID          string
	Hub         string
	Topic       string
	CallbackURL string
	CreatedAt   time.Time
	UpdatedAt   time.Time
	ExpiresAt   *time.Time
}

type PubSubState interface {
	All() (subscriptions []PubSubSubscription, err error)
	Add(hub, topic, baseURL string) (subscription *PubSubSubscription, oldCallbackURL string, err error)
	Get(hub, topic string) (subscription *PubSubSubscription, err error)
	GetByID(id string) (subscription *PubSubSubscription, err error)
	Set(hub, topic string, updatedAt, expiresAt time.Time) (err error)
	Del(hub, topic string) (err error)
}

type PubSubMessageHandler func(id string, s *PubSubSubscription, rd io.ReadCloser)

type PubSubClient struct {
	CallbackURL string
	State       PubSubState
	OnMessage   PubSubMessageHandler
}

func NewPubSubClient(callbackURL string, state PubSubState, onMessage PubSubMessageHandler) *PubSubClient {
	return &PubSubClient{
		CallbackURL: callbackURL,
		State:       state,
		OnMessage:   onMessage,
	}
}

func (c *PubSubClient) Refresh(forceUpdate bool, interval time.Duration) error {
	a, err := c.State.All()
	if err != nil {
		return errors.Wrap(err, "PubSubClient.Refresh")
	}

	logrus.WithField("count", len(a)).Debug("pubsub: got subscriptions to refresh")

	m := make(map[string]*ratelimit.Bucket)
	var l sync.RWMutex

	getBucket := func(host string) *ratelimit.Bucket {
		l.RLock()
		if b, ok := m[host]; ok {
			l.RUnlock()
			return b
		}
		l.RUnlock()

		l.Lock()
		if b, ok := m[host]; ok {
			l.Unlock()
			return b
		}
		defer l.Unlock()

		m[host] = ratelimit.NewBucket(time.Second*30, 4)

		return m[host]
	}

	var g WorkerGroup

	for _, e := range a {
		e := e

		callbackURL := c.CallbackURL + "/" + e.ID

		l := logrus.WithFields(logrus.Fields{
			"id":               e.ID,
			"hub":              e.Hub,
			"topic":            e.Topic,
			"callback_url":     e.CallbackURL,
			"new_callback_url": callbackURL,
			"created_at":       e.CreatedAt,
			"updated_at":       e.UpdatedAt,
			"expires_at":       e.ExpiresAt,
			"force_update":     forceUpdate,
		})

		if e.ExpiresAt == nil {
			l.Debug("pubsub: not refreshing subscription which has never been verified")
			continue
		}

		if forceUpdate || e.ExpiresAt.Sub(time.Now()) < interval || e.CallbackURL != callbackURL {
			u, err := url.Parse(e.Hub)
			if err != nil {
				l.WithError(err).Warn("pubsub: couldn't parse hub url")
				return errors.Wrap(err, "PubSubClient.RefreshWorker")
			}

			dur, ok := getBucket(u.Host).TakeMaxDuration(1, *pubsubRefreshInterval)
			if !ok {
				l.Debug("pubsub: skipping renewing for now as we'd have to wait too long")
				continue
			}

			g.Add(func() error {
				if dur > 0 {
					l.WithField("duration", dur).Debug("pubsub: waiting so as not to overwhelm the endpoint")
					time.Sleep(dur)
				}

				l.Debug("pubsub: refreshing subscription")

				// simulate subscription
				time.Sleep(time.Second * 5)

				// if err := c.Subscribe(e.Hub, e.Topic); err != nil {
				// 	l.WithError(err).Warn("pubsub: couldn't subscribe to topic")
				// 	return errors.Wrap(err, "PubSubClient.RefreshWorker")
				// }

				l.Debug("pubsub: subscribed successfully")
				return nil
			})
		}
	}

	return errors.Wrap(g.Run(20), "PubSubClient.Refresh")
}

func (c *PubSubClient) Subscribe(hub, topic string) error {
	logrus.WithFields(logrus.Fields{"hub": hub, "topic": topic}).Debug("pubsub: subscribing")

	s, oldCallbackURL, err := c.State.Add(hub, topic, c.CallbackURL)
	if err != nil {
		return errors.Wrap(err, "PubSubClient.Subscribe")
	}

	if oldCallbackURL != s.CallbackURL {
		if oldCallbackURL != "" {
			if err := PubSubUnsubscribe(hub, topic, oldCallbackURL); err != nil {
				return errors.Wrap(err, "PubSubClient.Subscribe")
			}
		}

		if err := PubSubSubscribe(hub, topic, s.CallbackURL); err != nil {
			return errors.Wrap(err, "PubSubClient.Subscribe")
		}
	}

	return nil
}

func (c *PubSubClient) Unsubscribe(hub, topic string) error {
	logrus.WithFields(logrus.Fields{"hub": hub, "topic": topic}).Debug("pubsub: unsubscribing")

	s, err := c.State.Get(hub, topic)
	if err != nil {
		return errors.Wrap(err, "PubSubClient.Unsubscribe")
	}

	if s != nil {
		if err := PubSubUnsubscribe(hub, topic, s.CallbackURL); err != nil {
			return errors.Wrap(err, "PubSubClient.Unsubscribe")
		}

		if err := c.State.Del(s.Hub, s.Topic); err != nil {
			return errors.Wrap(err, "PubSubClient.Unsubscribe")
		}
	}

	return nil
}

func (c *PubSubClient) Handler() *PubSubHandler {
	return &PubSubHandler{
		OnChallenge: func(id, topic, mode string, leaseTime time.Duration) error {
			l := logrus.WithFields(logrus.Fields{
				"id":         id,
				"topic":      topic,
				"mode":       mode,
				"lease_time": leaseTime,
			})

			l.Debug("pubsub: received challenge")

			s, err := c.State.GetByID(id)
			if err != nil {
				l.WithError(err).Warn("pubsub: error fetching subscription during challenge")
				return err
			}

			if s == nil {
				l.WithError(err).Warn("pubsub: subscription not found during challenge")
				return errors.Errorf("PubSubClient.Handler: subscription not found")
			}

			return errors.Wrap(c.State.Set(s.Hub, s.Topic, time.Now(), time.Now().Add(leaseTime)), "PubSubClient.Handler")
		},
		OnMessage: func(id, topic string, rd io.ReadCloser) {
			l := logrus.WithFields(logrus.Fields{
				"id":    id,
				"topic": topic,
			})

			l.Debug("pubsub: received message")

			defer rd.Close()

			if c.OnMessage != nil {
				s, err := c.State.GetByID(id)
				if err != nil {
					l.WithError(err).Warn("pubsub: error fetching subscription during reception")
					return
				}

				c.OnMessage(id, s, rd)
			}
		},
	}
}

func PubSubAlter(hub, topic, callbackURL, mode string) error {
	res, err := http.PostForm(hub, url.Values{
		"hub.callback":      []string{callbackURL},
		"hub.mode":          []string{mode},
		"hub.topic":         []string{topic},
		"hub.verify":        []string{"async"},
		"hub.lease_seconds": []string{"604800"},
	})

	if err != nil {
		return errors.Wrap(err, "PubSubAlter")
	}

	if res.StatusCode < 200 || res.StatusCode >= 300 {
		return errors.Errorf("PubSubAlter: invalid status code; expected 2xx but got %d", res.StatusCode)
	}

	return nil
}

func PubSubSubscribe(hub, topic, callbackURL string) error {
	return errors.Wrap(PubSubAlter(hub, topic, callbackURL, "subscribe"), "PubSubSubscribe")
}

func PubSubUnsubscribe(hub, topic, callbackURL string) error {
	return errors.Wrap(PubSubAlter(hub, topic, callbackURL, "unsubscribe"), "PubSubSubscribe")
}

type PubSubHandler struct {
	OnChallenge func(id, topic, mode string, leaseTime time.Duration) error
	OnMessage   func(id, topic string, rd io.ReadCloser)
}

func (h *PubSubHandler) ServeHTTP(rw http.ResponseWriter, r *http.Request) {
	q := r.URL.Query()

	bits := strings.Split(r.URL.Path, "/")
	id := bits[len(bits)-1]

	if r.Method == http.MethodGet && q.Get("hub.challenge") != "" {
		if h.OnChallenge != nil {
			leaseTime := time.Hour * 12
			if n, err := strconv.ParseInt(q.Get("hub.lease_seconds"), 10, 64); err == nil {
				leaseTime = time.Second * time.Duration(n)
			}

			if err := h.OnChallenge(id, q.Get("hub.topic"), q.Get("hub.mode"), leaseTime); err != nil {
				rw.WriteHeader(http.StatusNotFound)
				return
			}
		}

		rw.Write([]byte(r.URL.Query().Get("hub.challenge")))
		return
	}

	if r.Method != http.MethodPost {
		rw.WriteHeader(http.StatusMethodNotAllowed)
		return
	}

	links := linkheader.ParseMultiple(r.Header["Link"])

	var topic string
	if a := links.FilterByRel("self"); len(a) > 0 {
		topic = a[0].URL
	}

	rw.WriteHeader(http.StatusAccepted)

	if h.OnMessage != nil {
		h.OnMessage(id, topic, r.Body)
	}
}
<|endoftext|>


<commit_before>package s3gof3r

import (
	"bytes"
	"crypto/md5"
	"encoding/base64"
	"encoding/xml"
	"fmt"
	"hash"
	"io"
	"math"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"
)

// defined by amazon
const (
	minPartSize = 5 * mb
	maxPartSize = 5 * gb
	maxObjSize  = 5 * tb
	maxNPart    = 10000
	md5Header   = "content-md5"
)

type part struct {
	r   io.ReadSeeker
	len int64
	b   *bytes.Buffer

	// read by xml encoder
	PartNumber int
	ETag       string

	// Used for checksum of checksums on completion
	contentMd5 string
}

type putter struct {
	url         url.URL
	client      *http.Client
	b           *Bucket
	concurrency int
	nTry        int
	c           *Config

	bufsz      int64
	buf        *bytes.Buffer
	ch         chan *part
	part       int
	closed     bool
	err        error
	wg         sync.WaitGroup
	md5OfParts hash.Hash
	md5        hash.Hash
	ETag       string

	bp *bp

	makes    int
	UploadId string
	xml      struct {
		XMLName string `xml:"CompleteMultipartUpload"`
		Part    []*part
	}
}

type completeXml struct {
	ETag string
}

// Sends an S3 multipart upload initiation request.
// See http://docs.amazonwebservices.com/AmazonS3/latest/dev/mpuoverview.html.
// The initial request returns an UploadId that we use to identify
// subsequent PUT requests.
func newPutter(url url.URL, h http.Header, c *Config, b *Bucket) (p *putter, err error) {
	p = new(putter)
	p.url = url
	p.client = c.Client
	p.b = b
	p.concurrency = c.Concurrency
	p.nTry = c.NTry
	p.c = c
	p.bufsz = max64(minPartSize, c.PartSize)
	resp, err := p.retryRequest("POST", url.String()+"?uploads", nil, h)
	if err != nil {
		return nil, err
	}
	defer checkClose(resp.Body, &err)
	if resp.StatusCode != 200 {
		return nil, newRespError(resp)
	}
	err = xml.NewDecoder(resp.Body).Decode(p)
	if err != nil {
		return nil, err
	}
	p.ch = make(chan *part)
	for i := 0; i < p.concurrency; i++ {
		go p.worker()
	}
	p.md5OfParts = md5.New()
	p.md5 = md5.New()

	p.bp = newBufferPool(p.bufsz)

	return p, nil
}

func (p *putter) Write(b []byte) (int, error) {
	if p.closed {
		p.abort()
		return 0, syscall.EINVAL
	}
	if p.err != nil {
		p.abort()
		return 0, p.err
	}
	if p.buf == nil {
		p.buf = <-p.bp.get
		p.buf.Reset()
	}
	n, err := p.buf.Write(b)
	if err != nil {
		p.abort()
		return n, err
	}

	if int64(p.buf.Len()) >= p.bufsz {
		p.flush()
	}
	return n, nil
}

func (p *putter) flush() {
	p.wg.Add(1)
	p.part++
	b := *p.buf
	part := &part{bytes.NewReader(b.Bytes()), int64(b.Len()), p.buf, p.part, "", ""}
	var err error
	part.contentMd5, part.ETag, err = p.md5Content(part.r)
	if err != nil {
		p.err = err
	}

	p.xml.Part = append(p.xml.Part, part)
	p.ch <- part
	p.buf = nil
	// double buffer size every 1000 parts to
	// avoid exceeding the 10000-part AWS limit
	// while still reaching the 5 Terabyte max object size
	if p.part%1000 == 0 {
		p.bufsz = min64(p.bufsz*2, maxPartSize)
	}

}

func (p *putter) worker() {
	for part := range p.ch {
		p.retryPutPart(part)
	}
}

// Calls putPart up to nTry times to recover from transient errors.
func (p *putter) retryPutPart(part *part) {
	defer p.wg.Done()
	var err error
	for i := 0; i < p.nTry; i++ {
		time.Sleep(time.Duration(math.Exp2(float64(i))) * 100 * time.Millisecond) // exponential back-off
		err = p.putPart(part)
		if err == nil {
			p.bp.give <- part.b
			return
		}
		logger.debugPrintf("Error on attempt %d: Retrying part: %v, Error: %s", i, part, err)
	}
	p.err = err
}

// uploads a part, checking the etag against the calculated value
func (p *putter) putPart(part *part) error {
	v := url.Values{}
	v.Set("partNumber", strconv.Itoa(part.PartNumber))
	v.Set("uploadId", p.UploadId)
	if _, err := part.r.Seek(0, 0); err != nil { // move back to beginning, if retrying
		return err
	}
	req, err := http.NewRequest("PUT", p.url.String()+"?"+v.Encode(), part.r)
	if err != nil {
		return err
	}
	req.ContentLength = part.len
	req.Header.Set(md5Header, part.contentMd5)
	p.b.Sign(req)
	resp, err := p.client.Do(req)
	if err != nil {
		return err
	}
	defer checkClose(resp.Body, &err)
	if resp.StatusCode != 200 {
		return newRespError(resp)
	}
	s := resp.Header.Get("etag")
	s = s[1 : len(s)-1] // includes quote chars for some reason
	if part.ETag != s {
		return fmt.Errorf("Response etag does not match. Remote:%s Calculated:%s", s, p.ETag)
	}
	return nil
}

func (p *putter) Close() (err error) {
	if p.closed {
		p.abort()
		return syscall.EINVAL
	}
	if p.buf != nil {
		buf := *p.buf
		if buf.Len() > 0 {
			p.flush()
		}
	}
	p.wg.Wait()
	close(p.ch)
	p.closed = true
	p.bp.quit <- true

	if p.part == 0 {
		p.abort()
		return fmt.Errorf("0 bytes written")
	}
	if p.err != nil {
		p.abort()
		return p.err
	}
	logger.debugPrintln("buffer allocations:", p.bp.makes)
	// Complete Multipart upload
	body, err := xml.Marshal(p.xml)
	if err != nil {
		p.abort()
		return
	}
	b := bytes.NewReader(body)
	v := url.Values{}
	v.Set("uploadId", p.UploadId)
	resp, err := p.retryRequest("POST", p.url.String()+"?"+v.Encode(), b, nil)
	if err != nil {
		p.abort()
		return
	}
	defer checkClose(resp.Body, &err)
	if resp.StatusCode != 200 {
		p.abort()
		return newRespError(resp)
	}
	// Check md5 hash of concatenated part md5 hashes against ETag
	// more info: https://forums.aws.amazon.com/thread.jspa?messageID=456442&#456442
	calculatedMd5ofParts := fmt.Sprintf("%x", p.md5OfParts.Sum(nil))
	// Parse etag from body of response
	err = xml.NewDecoder(resp.Body).Decode(p)
	if err != nil {
		return
	}
	// strip part count from end and '"' from front.
	remoteMd5ofParts := strings.Split(p.ETag, "-")[0]
	remoteMd5ofParts = remoteMd5ofParts[1:len(remoteMd5ofParts)]
	if calculatedMd5ofParts != remoteMd5ofParts {
		if err != nil {
			return err
		}
		return fmt.Errorf("MD5 hash of part hashes comparison failed. Hash from multipart complete header: %s."+
			" Calculated multipart hash: %s.", remoteMd5ofParts, calculatedMd5ofParts)
	}
	if p.c.Md5Check {
		for i := 0; i < p.nTry; i++ {
			if err = p.putMd5(); err == nil {
				break
			}
		}
		return
	}
	return
}

// Try to abort multipart upload. Do not error on failure.
func (p *putter) abort() {
	v := url.Values{}
	v.Set("uploadId", p.UploadId)
	s := p.url.String() + "?" + v.Encode()
	resp, err := p.retryRequest("DELETE", s, nil, nil)
	if err != nil {
		logger.Printf("Error aborting multipart upload: %v", err)
		return
	}
	defer checkClose(resp.Body, &err)
	if resp.StatusCode != 204 {
		logger.Printf("Error aborting multipart upload: %v", newRespError(resp))
	}
	return
}

// Md5 functions
func (p *putter) md5Content(r io.ReadSeeker) (string, string, error) {
	h := md5.New()
	mw := io.MultiWriter(h, p.md5)
	if _, err := io.Copy(mw, r); err != nil {
		return "", "", err
	}
	sum := h.Sum(nil)
	hexSum := fmt.Sprintf("%x", sum)
	// add to checksum of all parts for verification on upload completion
	if _, err := p.md5OfParts.Write(sum); err != nil {
		return "", "", err
	}
	return base64.StdEncoding.EncodeToString(sum), hexSum, nil
}

// Put md5 file in .md5 subdirectory of bucket  where the file is stored
// e.g. the md5 for https://mybucket.s3.amazonaws.com/gof3r will be stored in
// https://mybucket.s3.amazonaws.com/.md5/gof3r.md5
func (p *putter) putMd5() (err error) {
	calcMd5 := fmt.Sprintf("%x", p.md5.Sum(nil))
	md5Reader := strings.NewReader(calcMd5)
	md5Path := fmt.Sprint(".md5", p.url.Path, ".md5")
	md5Url := p.b.Url(md5Path, p.c)
	logger.debugPrintln("md5: ", calcMd5)
	logger.debugPrintln("md5Path: ", md5Path)
	r, err := http.NewRequest("PUT", md5Url.String(), md5Reader)
	if err != nil {
		return
	}
	p.b.Sign(r)
	resp, err := p.client.Do(r)
	if err != nil {
		return
	}
	defer checkClose(resp.Body, &err)
	if resp.StatusCode != 200 {
		return newRespError(resp)
	}
	return
}

func (p *putter) retryRequest(method, urlStr string, body io.ReadSeeker, h http.Header) (resp *http.Response, err error) {
	for i := 0; i < p.nTry; i++ {
		var req *http.Request
		req, err = http.NewRequest(method, urlStr, body)
		if err != nil {
			return
		}
		for k := range h {
			for _, v := range h[k] {
				req.Header.Add(k, v)
			}
		}

		p.b.Sign(req)
		resp, err = p.client.Do(req)
		if err == nil {
			return
		}
		logger.debugPrintln(err)
		if body != nil {
			if _, err = body.Seek(0, 0); err != nil {
				return
			}
		}
	}
	return
}
<commit_msg>Remove redundant struct fields from putter.<commit_after>package s3gof3r

import (
	"bytes"
	"crypto/md5"
	"encoding/base64"
	"encoding/xml"
	"fmt"
	"hash"
	"io"
	"math"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"
)

// defined by amazon
const (
	minPartSize = 5 * mb
	maxPartSize = 5 * gb
	maxObjSize  = 5 * tb
	maxNPart    = 10000
	md5Header   = "content-md5"
)

type part struct {
	r   io.ReadSeeker
	len int64
	b   *bytes.Buffer

	// read by xml encoder
	PartNumber int
	ETag       string

	// Used for checksum of checksums on completion
	contentMd5 string
}

type putter struct {
	url url.URL
	b   *Bucket
	c   *Config

	bufsz      int64
	buf        *bytes.Buffer
	ch         chan *part
	part       int
	closed     bool
	err        error
	wg         sync.WaitGroup
	md5OfParts hash.Hash
	md5        hash.Hash
	ETag       string

	bp *bp

	makes    int
	UploadId string
	xml      struct {
		XMLName string `xml:"CompleteMultipartUpload"`
		Part    []*part
	}
}

type completeXml struct {
	ETag string
}

// Sends an S3 multipart upload initiation request.
// See http://docs.amazonwebservices.com/AmazonS3/latest/dev/mpuoverview.html.
// The initial request returns an UploadId that we use to identify
// subsequent PUT requests.
func newPutter(url url.URL, h http.Header, c *Config, b *Bucket) (p *putter, err error) {
	p = new(putter)
	p.url = url
	p.b = b
	p.c = c
	p.c.Concurrency = max(c.Concurrency, 1)
	p.c.NTry = max(c.NTry, 1)
	p.bufsz = max64(minPartSize, c.PartSize)
	resp, err := p.retryRequest("POST", url.String()+"?uploads", nil, h)
	if err != nil {
		return nil, err
	}
	defer checkClose(resp.Body, &err)
	if resp.StatusCode != 200 {
		return nil, newRespError(resp)
	}
	err = xml.NewDecoder(resp.Body).Decode(p)
	if err != nil {
		return nil, err
	}
	p.ch = make(chan *part)
	for i := 0; i < p.c.Concurrency; i++ {
		go p.worker()
	}
	p.md5OfParts = md5.New()
	p.md5 = md5.New()

	p.bp = newBufferPool(p.bufsz)

	return p, nil
}

func (p *putter) Write(b []byte) (int, error) {
	if p.closed {
		p.abort()
		return 0, syscall.EINVAL
	}
	if p.err != nil {
		p.abort()
		return 0, p.err
	}
	if p.buf == nil {
		p.buf = <-p.bp.get
		p.buf.Reset()
	}
	n, err := p.buf.Write(b)
	if err != nil {
		p.abort()
		return n, err
	}

	if int64(p.buf.Len()) >= p.bufsz {
		p.flush()
	}
	return n, nil
}

func (p *putter) flush() {
	p.wg.Add(1)
	p.part++
	b := *p.buf
	part := &part{bytes.NewReader(b.Bytes()), int64(b.Len()), p.buf, p.part, "", ""}
	var err error
	part.contentMd5, part.ETag, err = p.md5Content(part.r)
	if err != nil {
		p.err = err
	}

	p.xml.Part = append(p.xml.Part, part)
	p.ch <- part
	p.buf = nil
	// double buffer size every 1000 parts to
	// avoid exceeding the 10000-part AWS limit
	// while still reaching the 5 Terabyte max object size
	if p.part%1000 == 0 {
		p.bufsz = min64(p.bufsz*2, maxPartSize)
	}

}

func (p *putter) worker() {
	for part := range p.ch {
		p.retryPutPart(part)
	}
}

// Calls putPart up to nTry times to recover from transient errors.
func (p *putter) retryPutPart(part *part) {
	defer p.wg.Done()
	var err error
	for i := 0; i < p.c.NTry; i++ {
		time.Sleep(time.Duration(math.Exp2(float64(i))) * 100 * time.Millisecond) // exponential back-off
		err = p.putPart(part)
		if err == nil {
			p.bp.give <- part.b
			return
		}
		logger.debugPrintf("Error on attempt %d: Retrying part: %v, Error: %s", i, part, err)
	}
	p.err = err
}

// uploads a part, checking the etag against the calculated value
func (p *putter) putPart(part *part) error {
	v := url.Values{}
	v.Set("partNumber", strconv.Itoa(part.PartNumber))
	v.Set("uploadId", p.UploadId)
	if _, err := part.r.Seek(0, 0); err != nil { // move back to beginning, if retrying
		return err
	}
	req, err := http.NewRequest("PUT", p.url.String()+"?"+v.Encode(), part.r)
	if err != nil {
		return err
	}
	req.ContentLength = part.len
	req.Header.Set(md5Header, part.contentMd5)
	p.b.Sign(req)
	resp, err := p.c.Client.Do(req)
	if err != nil {
		return err
	}
	defer checkClose(resp.Body, &err)
	if resp.StatusCode != 200 {
		return newRespError(resp)
	}
	s := resp.Header.Get("etag")
	s = s[1 : len(s)-1] // includes quote chars for some reason
	if part.ETag != s {
		return fmt.Errorf("Response etag does not match. Remote:%s Calculated:%s", s, p.ETag)
	}
	return nil
}

func (p *putter) Close() (err error) {
	if p.closed {
		p.abort()
		return syscall.EINVAL
	}
	if p.buf != nil {
		buf := *p.buf
		if buf.Len() > 0 {
			p.flush()
		}
	}
	p.wg.Wait()
	close(p.ch)
	p.closed = true
	p.bp.quit <- true

	if p.part == 0 {
		p.abort()
		return fmt.Errorf("0 bytes written")
	}
	if p.err != nil {
		p.abort()
		return p.err
	}
	logger.debugPrintln("buffer allocations:", p.bp.makes)
	// Complete Multipart upload
	body, err := xml.Marshal(p.xml)
	if err != nil {
		p.abort()
		return
	}
	b := bytes.NewReader(body)
	v := url.Values{}
	v.Set("uploadId", p.UploadId)
	resp, err := p.retryRequest("POST", p.url.String()+"?"+v.Encode(), b, nil)
	if err != nil {
		p.abort()
		return
	}
	defer checkClose(resp.Body, &err)
	if resp.StatusCode != 200 {
		p.abort()
		return newRespError(resp)
	}
	// Check md5 hash of concatenated part md5 hashes against ETag
	// more info: https://forums.aws.amazon.com/thread.jspa?messageID=456442&#456442
	calculatedMd5ofParts := fmt.Sprintf("%x", p.md5OfParts.Sum(nil))
	// Parse etag from body of response
	err = xml.NewDecoder(resp.Body).Decode(p)
	if err != nil {
		return
	}
	// strip part count from end and '"' from front.
	remoteMd5ofParts := strings.Split(p.ETag, "-")[0]
	remoteMd5ofParts = remoteMd5ofParts[1:len(remoteMd5ofParts)]
	if calculatedMd5ofParts != remoteMd5ofParts {
		if err != nil {
			return err
		}
		return fmt.Errorf("MD5 hash of part hashes comparison failed. Hash from multipart complete header: %s."+
			" Calculated multipart hash: %s.", remoteMd5ofParts, calculatedMd5ofParts)
	}
	if p.c.Md5Check {
		for i := 0; i < p.c.NTry; i++ {
			if err = p.putMd5(); err == nil {
				break
			}
		}
		return
	}
	return
}

// Try to abort multipart upload. Do not error on failure.
func (p *putter) abort() {
	v := url.Values{}
	v.Set("uploadId", p.UploadId)
	s := p.url.String() + "?" + v.Encode()
	resp, err := p.retryRequest("DELETE", s, nil, nil)
	if err != nil {
		logger.Printf("Error aborting multipart upload: %v\n", err)
		return
	}
	defer checkClose(resp.Body, &err)
	if resp.StatusCode != 204 {
		logger.Printf("Error aborting multipart upload: %v", newRespError(resp))
	}
	return
}

// Md5 functions
func (p *putter) md5Content(r io.ReadSeeker) (string, string, error) {
	h := md5.New()
	mw := io.MultiWriter(h, p.md5)
	if _, err := io.Copy(mw, r); err != nil {
		return "", "", err
	}
	sum := h.Sum(nil)
	hexSum := fmt.Sprintf("%x", sum)
	// add to checksum of all parts for verification on upload completion
	if _, err := p.md5OfParts.Write(sum); err != nil {
		return "", "", err
	}
	return base64.StdEncoding.EncodeToString(sum), hexSum, nil
}

// Put md5 file in .md5 subdirectory of bucket  where the file is stored
// e.g. the md5 for https://mybucket.s3.amazonaws.com/gof3r will be stored in
// https://mybucket.s3.amazonaws.com/.md5/gof3r.md5
func (p *putter) putMd5() (err error) {
	calcMd5 := fmt.Sprintf("%x", p.md5.Sum(nil))
	md5Reader := strings.NewReader(calcMd5)
	md5Path := fmt.Sprint(".md5", p.url.Path, ".md5")
	md5Url := p.b.Url(md5Path, p.c)
	logger.debugPrintln("md5: ", calcMd5)
	logger.debugPrintln("md5Path: ", md5Path)
	r, err := http.NewRequest("PUT", md5Url.String(), md5Reader)
	if err != nil {
		return
	}
	p.b.Sign(r)
	resp, err := p.c.Client.Do(r)
	if err != nil {
		return
	}
	defer checkClose(resp.Body, &err)
	if resp.StatusCode != 200 {
		return newRespError(resp)
	}
	return
}

func (p *putter) retryRequest(method, urlStr string, body io.ReadSeeker, h http.Header) (resp *http.Response, err error) {
	for i := 0; i < p.c.NTry; i++ {
		var req *http.Request
		req, err = http.NewRequest(method, urlStr, body)
		if err != nil {
			return
		}
		for k := range h {
			for _, v := range h[k] {
				req.Header.Add(k, v)
			}
		}

		p.b.Sign(req)
		resp, err = p.c.Client.Do(req)
		if err == nil {
			return
		}
		logger.debugPrintln(err)
		if body != nil {
			if _, err = body.Seek(0, 0); err != nil {
				return
			}
		}
	}
	return
}
<|endoftext|>


<commit_before>package jsonapi

import (
	"fmt"
	"io"
	"reflect"
	"runtime"

	"github.com/nstratos/jsonapi"
)

func isZeroOfUnderlyingType(x interface{}) bool {
	return reflect.DeepEqual(x, reflect.Zero(reflect.TypeOf(x)).Interface())
}

func Encode(w io.Writer, v interface{}) (err error) {
	const errFormat = "cannot encode type %T, need pointer to struct or slice of pointers to structs"
	defer func() {
		if r := recover(); r != nil {
			if _, ok := r.(runtime.Error); ok {
				panic(r)
			}
			err = fmt.Errorf(errFormat+": %v", v, r.(error))
		}
	}()
	if isZeroOfUnderlyingType(v) {
		return fmt.Errorf("cannot encode nil value of %#v", v)
	}
	t := reflect.TypeOf(v)
	switch t.Kind() {
	default:
		return fmt.Errorf(errFormat, v)
	case reflect.Ptr:
		if t.Elem().Kind() != reflect.Struct {
			return fmt.Errorf(errFormat, v)
		}
		return jsonapi.MarshalOnePayload(w, v)
	case reflect.Slice:
		s := reflect.ValueOf(v)
		if s.Type().Elem().Kind() != reflect.Ptr {
			return fmt.Errorf(errFormat, v)
		}
		if s.Type().Elem().Elem().Kind() != reflect.Struct {
			return fmt.Errorf(errFormat, v)

		}
		return jsonapi.MarshalManyPayload(w, v)
	}
}

func Decode(r io.Reader, ptr interface{}) (offset Offset, err error) {
	const errFormat = "cannot decode to %T, need pointer to struct or pointer to slice"
	defer func() {
		if r := recover(); r != nil {
			if _, ok := r.(runtime.Error); ok {
				panic(r)
			}
			err = fmt.Errorf(errFormat+": %v", ptr, r)
		}
	}()
	if reflect.TypeOf(ptr).Kind() != reflect.Ptr {
		return Offset{}, fmt.Errorf(errFormat, ptr)
	}
	v := reflect.Indirect(reflect.ValueOf(ptr))
	switch v.Kind() {
	default:
		return Offset{}, fmt.Errorf(errFormat, ptr)
	case reflect.Struct:
		return Offset{}, jsonapi.UnmarshalPayload(r, ptr)
	case reflect.Slice:
		data, links, uerr := jsonapi.UnmarshalManyPayloadWithLinks(r, v.Type().Elem())
		if uerr != nil {
			return Offset{}, uerr
		}
		for _, d := range data {
			v.Set(reflect.Append(v, reflect.ValueOf(d)))
		}

		o := Offset{}
		var perr error
		if links != nil {
			o, perr = parseOffset(*links)
			if perr != nil {
				return Offset{}, perr
			}
		}
		return o, nil
	}
}

func EncodeOne(w io.Writer, v interface{}) error {
	return jsonapi.MarshalOnePayload(w, v)
}

func EncodeMany(w io.Writer, v interface{}) error {
	return jsonapi.MarshalManyPayload(w, v)
}

func DecodeOne(r io.Reader, v interface{}) error {
	return jsonapi.UnmarshalPayload(r, v)
}

func DecodeMany(r io.Reader, t reflect.Type) ([]interface{}, Offset, error) {
	v, links, err := jsonapi.UnmarshalManyPayloadWithLinks(r, t)
	if err != nil {
		return nil, Offset{}, err
	}

	o := Offset{}
	var perr error
	if links != nil {
		o, perr = parseOffset(*links)
		if perr != nil {
			return nil, Offset{}, perr
		}
	}
	return v, o, err
}
<commit_msg>Remove type assertion in recover and just print it<commit_after>package jsonapi

import (
	"fmt"
	"io"
	"reflect"
	"runtime"

	"github.com/nstratos/jsonapi"
)

func isZeroOfUnderlyingType(x interface{}) bool {
	return reflect.DeepEqual(x, reflect.Zero(reflect.TypeOf(x)).Interface())
}

func Encode(w io.Writer, v interface{}) (err error) {
	const errFormat = "cannot encode type %T, need pointer to struct or slice of pointers to structs"
	defer func() {
		if r := recover(); r != nil {
			if _, ok := r.(runtime.Error); ok {
				panic(r)
			}
			err = fmt.Errorf(errFormat+": %v", v, r)
		}
	}()
	if isZeroOfUnderlyingType(v) {
		return fmt.Errorf("cannot encode nil value of %#v", v)
	}
	t := reflect.TypeOf(v)
	switch t.Kind() {
	default:
		return fmt.Errorf(errFormat, v)
	case reflect.Ptr:
		if t.Elem().Kind() != reflect.Struct {
			return fmt.Errorf(errFormat, v)
		}
		return jsonapi.MarshalOnePayload(w, v)
	case reflect.Slice:
		s := reflect.ValueOf(v)
		if s.Type().Elem().Kind() != reflect.Ptr {
			return fmt.Errorf(errFormat, v)
		}
		if s.Type().Elem().Elem().Kind() != reflect.Struct {
			return fmt.Errorf(errFormat, v)

		}
		return jsonapi.MarshalManyPayload(w, v)
	}
}

func Decode(r io.Reader, ptr interface{}) (offset Offset, err error) {
	const errFormat = "cannot decode to %T, need pointer to struct or pointer to slice"
	defer func() {
		if r := recover(); r != nil {
			if _, ok := r.(runtime.Error); ok {
				panic(r)
			}
			err = fmt.Errorf(errFormat+": %v", ptr, r)
		}
	}()
	if reflect.TypeOf(ptr).Kind() != reflect.Ptr {
		return Offset{}, fmt.Errorf(errFormat, ptr)
	}
	v := reflect.Indirect(reflect.ValueOf(ptr))
	switch v.Kind() {
	default:
		return Offset{}, fmt.Errorf(errFormat, ptr)
	case reflect.Struct:
		return Offset{}, jsonapi.UnmarshalPayload(r, ptr)
	case reflect.Slice:
		data, links, uerr := jsonapi.UnmarshalManyPayloadWithLinks(r, v.Type().Elem())
		if uerr != nil {
			return Offset{}, uerr
		}
		for _, d := range data {
			v.Set(reflect.Append(v, reflect.ValueOf(d)))
		}

		o := Offset{}
		var perr error
		if links != nil {
			o, perr = parseOffset(*links)
			if perr != nil {
				return Offset{}, perr
			}
		}
		return o, nil
	}
}

func EncodeOne(w io.Writer, v interface{}) error {
	return jsonapi.MarshalOnePayload(w, v)
}

func EncodeMany(w io.Writer, v interface{}) error {
	return jsonapi.MarshalManyPayload(w, v)
}

func DecodeOne(r io.Reader, v interface{}) error {
	return jsonapi.UnmarshalPayload(r, v)
}

func DecodeMany(r io.Reader, t reflect.Type) ([]interface{}, Offset, error) {
	v, links, err := jsonapi.UnmarshalManyPayloadWithLinks(r, t)
	if err != nil {
		return nil, Offset{}, err
	}

	o := Offset{}
	var perr error
	if links != nil {
		o, perr = parseOffset(*links)
		if perr != nil {
			return nil, Offset{}, perr
		}
	}
	return v, o, err
}
<|endoftext|>


<commit_before>package cs

import (
	"github.com/denverdino/aliyungo/common"
	"net/http"
	"net/url"
	"strconv"
	"time"
)

type ServerlessCreationArgs struct {
	ClusterType          KubernetesClusterType `json:"cluster_type"`
	Name                 string                `json:"name"`
	RegionId             string                `json:"region_id"`
	VpcId                string                `json:"vpc_id"`
	VSwitchId            string                `json:"vswitch_id"`
	VswitchIds           []string              `json:"vswitch_ids"`
	EndpointPublicAccess bool                  `json:"public_slb"`
	PrivateZone          bool                  `json:"private_zone"`
	NatGateway           bool                  `json:"nat_gateway"`
	DeletionProtection   bool                  `json:"deletion_protection"`
	SecurityGroupId      string                `json:"security_group_id"`
	Tags                 []Tag                 `json:"tags"`
	Addons               []Addon               `json:"addons"`
}

type ServerlessClusterResponse struct {
	ClusterId          string                `json:"cluster_id"`
	Name               string                `json:"name"`
	ClusterType        KubernetesClusterType `json:"cluster_type"`
	RegionId           string                `json:"region_id"`
	State              ClusterState          `json:"state"`
	VpcId              string                `json:"vpc_id"`
	VSwitchId          string                `json:"vswitch_id"`
	SecurityGroupId    string                `json:"security_group_id"`
	Tags               []Tag                 `json:"tags"`
	Created            time.Time             `json:"created"`
	Updated            time.Time             `json:"updated"`
	InitVersion        string                `json:"init_version"`
	CurrentVersion     string                `json:"current_version"`
	PrivateZone        bool                  `json:"private_zone"`
	DeletionProtection bool                  `json:"deletion_protection"`
}

type Tag struct {
	Key   string `json:"key"`
	Value string `json:"value"`
}

func (this *ServerlessCreationArgs) Validate() error {
	if this.Name == "" || this.RegionId == "" || this.VpcId == "" || this.VSwitchId == "" {
		return common.GetCustomError("InvalidParameters", "The name,region_id,vpc_id,vswitch_id not allowed empty")
	}
	return nil
}

//create Serverless cluster
func (client *Client) CreateServerlessKubernetesCluster(args *ServerlessCreationArgs) (*ClusterCommonResponse, error) {
	if args == nil {
		return nil, common.GetCustomError("InvalidArgs", "The args is nil ")
	}

	if err := args.Validate(); err != nil {
		return nil, err
	}

	//reset clusterType,
	args.ClusterType = ServerlessKubernetes
	cluster := &ClusterCommonResponse{}
	err := client.Invoke(common.Region(args.RegionId), http.MethodPost, "/clusters", nil, args, &cluster)
	if err != nil {
		return nil, err
	}
	return cluster, nil
}

//describe Serverless cluster
func (client *Client) DescribeServerlessKubernetesCluster(clusterId string) (*ServerlessClusterResponse, error) {
	cluster := &ServerlessClusterResponse{}
	err := client.Invoke("", http.MethodGet, "/clusters/"+clusterId, nil, nil, cluster)
	if err != nil {
		return nil, err
	}
	return cluster, nil
}

//describe Serverless clsuters
func (client *Client) DescribeServerlessKubernetesClusters() ([]*ServerlessClusterResponse, error) {
	allClusters := make([]*ServerlessClusterResponse, 0)
	askClusters := make([]*ServerlessClusterResponse, 0)

	err := client.Invoke("", http.MethodGet, "/clusters", nil, nil, &allClusters)
	if err != nil {
		return askClusters, err
	}

	for _, cluster := range allClusters {
		if cluster.ClusterType == ClusterTypeServerlessKubernetes {
			askClusters = append(askClusters, cluster)
		}
	}

	return askClusters, nil
}

//new api for get cluster kube user config
func (client *Client) DescribeClusterUserConfig(clusterId string, privateIpAddress bool) (*ClusterConfig, error) {
	config := &ClusterConfig{}
	query := url.Values{}
	query.Add("PrivateIpAddress", strconv.FormatBool(privateIpAddress))

	err := client.Invoke("", http.MethodGet, "/k8s/"+clusterId+"/user_config", query, nil, &config)
	if err != nil {
		return nil, err
	}
	return config, nil
}
<commit_msg>create ask support vswitch_ids<commit_after>package cs

import (
	"github.com/denverdino/aliyungo/common"
	"net/http"
	"net/url"
	"strconv"
	"time"
)

type ServerlessCreationArgs struct {
	ClusterType          KubernetesClusterType `json:"cluster_type"`
	Name                 string                `json:"name"`
	RegionId             string                `json:"region_id"`
	VpcId                string                `json:"vpc_id"`
	VSwitchId            string                `json:"vswitch_id"`
	VswitchIds           []string              `json:"vswitch_ids"`
	EndpointPublicAccess bool                  `json:"public_slb"`
	PrivateZone          bool                  `json:"private_zone"`
	NatGateway           bool                  `json:"nat_gateway"`
	DeletionProtection   bool                  `json:"deletion_protection"`
	SecurityGroupId      string                `json:"security_group_id"`
	Tags                 []Tag                 `json:"tags"`
	Addons               []Addon               `json:"addons"`
}

type ServerlessClusterResponse struct {
	ClusterId          string                `json:"cluster_id"`
	Name               string                `json:"name"`
	ClusterType        KubernetesClusterType `json:"cluster_type"`
	RegionId           string                `json:"region_id"`
	State              ClusterState          `json:"state"`
	VpcId              string                `json:"vpc_id"`
	VSwitchId          string                `json:"vswitch_id"`
	SecurityGroupId    string                `json:"security_group_id"`
	Tags               []Tag                 `json:"tags"`
	Created            time.Time             `json:"created"`
	Updated            time.Time             `json:"updated"`
	InitVersion        string                `json:"init_version"`
	CurrentVersion     string                `json:"current_version"`
	PrivateZone        bool                  `json:"private_zone"`
	DeletionProtection bool                  `json:"deletion_protection"`
}

type Tag struct {
	Key   string `json:"key"`
	Value string `json:"value"`
}

func (this *ServerlessCreationArgs) Validate() error {
	if this.Name == "" || this.RegionId == "" || this.VpcId == "" {
		return common.GetCustomError("InvalidParameters", "The name,region_id,vpc_id not allowed empty")
	}
	return nil
}

//create Serverless cluster
func (client *Client) CreateServerlessKubernetesCluster(args *ServerlessCreationArgs) (*ClusterCommonResponse, error) {
	if args == nil {
		return nil, common.GetCustomError("InvalidArgs", "The args is nil ")
	}

	if err := args.Validate(); err != nil {
		return nil, err
	}

	//reset clusterType,
	args.ClusterType = ServerlessKubernetes
	cluster := &ClusterCommonResponse{}
	err := client.Invoke(common.Region(args.RegionId), http.MethodPost, "/clusters", nil, args, &cluster)
	if err != nil {
		return nil, err
	}
	return cluster, nil
}

//describe Serverless cluster
func (client *Client) DescribeServerlessKubernetesCluster(clusterId string) (*ServerlessClusterResponse, error) {
	cluster := &ServerlessClusterResponse{}
	err := client.Invoke("", http.MethodGet, "/clusters/"+clusterId, nil, nil, cluster)
	if err != nil {
		return nil, err
	}
	return cluster, nil
}

//describe Serverless clsuters
func (client *Client) DescribeServerlessKubernetesClusters() ([]*ServerlessClusterResponse, error) {
	allClusters := make([]*ServerlessClusterResponse, 0)
	askClusters := make([]*ServerlessClusterResponse, 0)

	err := client.Invoke("", http.MethodGet, "/clusters", nil, nil, &allClusters)
	if err != nil {
		return askClusters, err
	}

	for _, cluster := range allClusters {
		if cluster.ClusterType == ClusterTypeServerlessKubernetes {
			askClusters = append(askClusters, cluster)
		}
	}

	return askClusters, nil
}

//new api for get cluster kube user config
func (client *Client) DescribeClusterUserConfig(clusterId string, privateIpAddress bool) (*ClusterConfig, error) {
	config := &ClusterConfig{}
	query := url.Values{}
	query.Add("PrivateIpAddress", strconv.FormatBool(privateIpAddress))

	err := client.Invoke("", http.MethodGet, "/k8s/"+clusterId+"/user_config", query, nil, &config)
	if err != nil {
		return nil, err
	}
	return config, nil
}
<|endoftext|>