<commit_before>// Copyright © 2016 The Things Network
// Use of this source code is governed by the MIT license that can be found in the LICENSE file.

package otaa

import (
	"crypto/aes"

	"github.com/TheThingsNetwork/ttn/core/types"
	"github.com/TheThingsNetwork/ttn/utils/errors"
)

// CalculateSessionKeys calculates the AppSKey and NwkSKey
// All arguments are MSB-first
func CalculateSessionKeys(appKey types.AppKey, appNonce [3]byte, netID [3]byte, devNonce [2]byte) (appSKey types.AppSKey, nwkSKey types.NwkSKey, err error) {

	buf := make([]byte, 16)
	copy(buf[1:4], reverse(appNonce[:]))
	copy(buf[4:7], reverse(netID[:]))
	copy(buf[7:9], reverse(devNonce[:]))

	block, err := aes.NewCipher(appKey[:])

	if err != nil || block.BlockSize() != 16 {
		err = errors.NewErrInternal("Unable to create cipher to generate keys")
		return
	}

	buf[0] = 0x1
	block.Encrypt(nwkSKey[:], buf)
	buf[0] = 0x2
	block.Encrypt(appSKey[:], buf)

	return
}

// reverse is used to convert between MSB-first and LSB-first
func reverse(in []byte) (out []byte) {
	for i := len(in) - 1; i >= 0; i-- {
		out = append(out, in[i])
	}
	return
}
<commit_msg>types.AppKey is always 16 bytes<commit_after>// Copyright © 2016 The Things Network
// Use of this source code is governed by the MIT license that can be found in the LICENSE file.

package otaa

import (
	"crypto/aes"

	"github.com/TheThingsNetwork/ttn/core/types"
)

// CalculateSessionKeys calculates the AppSKey and NwkSKey
// All arguments are MSB-first
func CalculateSessionKeys(appKey types.AppKey, appNonce [3]byte, netID [3]byte, devNonce [2]byte) (appSKey types.AppSKey, nwkSKey types.NwkSKey, err error) {

	buf := make([]byte, 16)
	copy(buf[1:4], reverse(appNonce[:]))
	copy(buf[4:7], reverse(netID[:]))
	copy(buf[7:9], reverse(devNonce[:]))

	block, _ := aes.NewCipher(appKey[:])

	buf[0] = 0x1
	block.Encrypt(nwkSKey[:], buf)
	buf[0] = 0x2
	block.Encrypt(appSKey[:], buf)

	return
}

// reverse is used to convert between MSB-first and LSB-first
func reverse(in []byte) (out []byte) {
	for i := len(in) - 1; i >= 0; i-- {
		out = append(out, in[i])
	}
	return
}
<|endoftext|>


<commit_before>// (c) 2019-2020, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package platformvm

import (
	"github.com/ava-labs/gecko/database"
	"github.com/ava-labs/gecko/database/versiondb"
	"github.com/ava-labs/gecko/ids"
	"github.com/ava-labs/gecko/snow/choices"
	"github.com/ava-labs/gecko/snow/consensus/snowman"
	"github.com/ava-labs/gecko/vms/components/core"
)

// ProposalTx is an operation that can be proposed
type ProposalTx interface {
	initialize(vm *VM) error
	// Attempts to verify this transaction with the provided state.
	SemanticVerify(database.Database) (
		onCommitDB *versiondb.Database,
		onAbortDB *versiondb.Database,
		onCommitFunc func(),
		onAbortFunc func(),
		err TxError,
	)
	InitiallyPrefersCommit() bool
}

// ProposalBlock is a proposal to change the chain's state.
// A proposal may be to:
// 	1. Advance the chain's timestamp (*AdvanceTimeTx)
//  2. Remove a staker from the staker set (*RewardStakerTx)
//  3. Add a new staker to the set of pending (future) stakers (*AddStakerTx)
// The proposal will be enacted (change the chain's state) if the proposal block
// is accepted and followed by an accepted Commit block
type ProposalBlock struct {
	CommonBlock `serialize:"true"`

	Tx ProposalTx `serialize:"true"`

	// The database that the chain will have if this block's proposal is committed
	onCommitDB *versiondb.Database
	// The database that the chain will have if this block's proposal is aborted
	onAbortDB *versiondb.Database
	// The function to execute if this block's proposal is committed
	onCommitFunc func()
	// The function to execute if this block's proposal is aborted
	onAbortFunc func()
}

// Accept implements the snowman.Block interface
func (pb *ProposalBlock) Accept() error {
	pb.SetStatus(choices.Accepted)
	pb.VM.LastAcceptedID = pb.ID()
	return nil
}

// Initialize this block.
// Sets [pb.vm] to [vm] and populates non-serialized fields
// This method should be called when a block is unmarshaled from bytes
func (pb *ProposalBlock) initialize(vm *VM, bytes []byte) error {
	pb.vm = vm
	pb.Block.Initialize(bytes, vm.SnowmanVM)
	return pb.Tx.initialize(vm)
}

// setBaseDatabase sets this block's base database to [db]
func (pb *ProposalBlock) setBaseDatabase(db database.Database) {
	if err := pb.onCommitDB.SetDatabase(db); err != nil {
		pb.vm.Ctx.Log.Error("problem while setting base database: %s", err)
	}
	if err := pb.onAbortDB.SetDatabase(db); err != nil {
		pb.vm.Ctx.Log.Error("problem while setting base database: %s", err)
	}
}

// onCommit should only be called after Verify is called.
// onCommit returns:
//   1. A database that contains the state of the chain assuming this proposal
//      is enacted. (That is, if this block is accepted and followed by an
//      accepted Commit block.)
//   2. A function be be executed when this block's proposal is committed.
//      This function should not write to state.
func (pb *ProposalBlock) onCommit() (*versiondb.Database, func()) {
	return pb.onCommitDB, pb.onCommitFunc
}

// onAbort should only be called after Verify is called.
// onAbort returns a database that contains the state of the chain assuming this
// block's proposal is rejected. (That is, if this block is accepted and
// followed by an accepted Abort block.)
func (pb *ProposalBlock) onAbort() (*versiondb.Database, func()) {
	return pb.onAbortDB, pb.onAbortFunc
}

// Verify this block is valid.
//
// The parent block must either be a Commit or an Abort block.
//
// If this block is valid, this function also sets pas.onCommit and pas.onAbort.
func (pb *ProposalBlock) Verify() error {
	parentIntf := pb.parentBlock()

	// The parent of a proposal block (ie this block) must be a decision block
	parent, ok := parentIntf.(decision)
	if !ok {
		if err := pb.Reject(); err == nil {
			if err := pb.vm.DB.Commit(); err != nil {
				pb.vm.Ctx.Log.Error("error committing Proposal block as rejected: %s", err)
			}
		} else {
			pb.vm.DB.Abort()
		}
		return errInvalidBlockType
	}

	// pdb is the database if this block's parent is accepted
	pdb := parent.onAccept()

	var err TxError
	pb.onCommitDB, pb.onAbortDB, pb.onCommitFunc, pb.onAbortFunc, err = pb.Tx.SemanticVerify(pdb)
	if err != nil {
		// If this block's transaction proposes to advance the timestamp, the transaction may fail
		// verification now but be valid in the future, so don't (permanently) mark the block as rejected.
		if !err.Temporary() {
			if err := pb.Reject(); err == nil {
				if err := pb.vm.DB.Commit(); err != nil {
					pb.vm.Ctx.Log.Error("error committing Proposal block as rejected: %s", err)
				}
			} else {
				pb.vm.DB.Abort()
			}
		}
		return err
	}

	pb.vm.currentBlocks[pb.ID().Key()] = pb
	parentIntf.addChild(pb)
	return nil
}

// Options returns the possible children of this block in preferential order.
func (pb *ProposalBlock) Options() ([2]snowman.Block, error) {
	blockID := pb.ID()

	commit, err := pb.vm.newCommitBlock(blockID)
	if err != nil {
		return [2]snowman.Block{}, err
	}
	abort, err := pb.vm.newAbortBlock(blockID)
	if err != nil {
		return [2]snowman.Block{}, err
	}

	if err := pb.vm.State.PutBlock(pb.vm.DB, commit); err != nil {
		pb.vm.Ctx.Log.Warn(errDBPutBlock.Error())
		return [2]snowman.Block{}, err
	}
	if err := pb.vm.State.PutBlock(pb.vm.DB, abort); err != nil {
		pb.vm.Ctx.Log.Warn(errDBPutBlock.Error())
		return [2]snowman.Block{}, err
	}
	pb.vm.DB.Commit()

	if pb.Tx.InitiallyPrefersCommit() {
		return [2]snowman.Block{commit, abort}, nil
	}
	return [2]snowman.Block{abort, commit}, nil
}

// newProposalBlock creates a new block that proposes to issue a transaction.
// The parent of this block has ID [parentID]. The parent must be a decision block.
// Returns nil if there's an error while creating this block
func (vm *VM) newProposalBlock(parentID ids.ID, tx ProposalTx) (*ProposalBlock, error) {
	pb := &ProposalBlock{
		CommonBlock: CommonBlock{
			Block: core.NewBlock(parentID),
			vm:    vm,
		},
		Tx: tx,
	}

	// We marshal the block in this way (as a Block) so that we can unmarshal
	// it into a Block (rather than a *ProposalBlock)
	block := Block(pb)
	bytes, err := Codec.Marshal(&block)
	if err != nil {
		return nil, err
	}
	pb.Initialize(bytes, vm.SnowmanVM)
	return pb, nil
}
<commit_msg>Added error handling for commital in the proposal_block Options function<commit_after>// (c) 2019-2020, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package platformvm

import (
	"github.com/ava-labs/gecko/database"
	"github.com/ava-labs/gecko/database/versiondb"
	"github.com/ava-labs/gecko/ids"
	"github.com/ava-labs/gecko/snow/choices"
	"github.com/ava-labs/gecko/snow/consensus/snowman"
	"github.com/ava-labs/gecko/vms/components/core"
)

// ProposalTx is an operation that can be proposed
type ProposalTx interface {
	initialize(vm *VM) error
	// Attempts to verify this transaction with the provided state.
	SemanticVerify(database.Database) (
		onCommitDB *versiondb.Database,
		onAbortDB *versiondb.Database,
		onCommitFunc func(),
		onAbortFunc func(),
		err TxError,
	)
	InitiallyPrefersCommit() bool
}

// ProposalBlock is a proposal to change the chain's state.
// A proposal may be to:
// 	1. Advance the chain's timestamp (*AdvanceTimeTx)
//  2. Remove a staker from the staker set (*RewardStakerTx)
//  3. Add a new staker to the set of pending (future) stakers (*AddStakerTx)
// The proposal will be enacted (change the chain's state) if the proposal block
// is accepted and followed by an accepted Commit block
type ProposalBlock struct {
	CommonBlock `serialize:"true"`

	Tx ProposalTx `serialize:"true"`

	// The database that the chain will have if this block's proposal is committed
	onCommitDB *versiondb.Database
	// The database that the chain will have if this block's proposal is aborted
	onAbortDB *versiondb.Database
	// The function to execute if this block's proposal is committed
	onCommitFunc func()
	// The function to execute if this block's proposal is aborted
	onAbortFunc func()
}

// Accept implements the snowman.Block interface
func (pb *ProposalBlock) Accept() error {
	pb.SetStatus(choices.Accepted)
	pb.VM.LastAcceptedID = pb.ID()
	return nil
}

// Initialize this block.
// Sets [pb.vm] to [vm] and populates non-serialized fields
// This method should be called when a block is unmarshaled from bytes
func (pb *ProposalBlock) initialize(vm *VM, bytes []byte) error {
	pb.vm = vm
	pb.Block.Initialize(bytes, vm.SnowmanVM)
	return pb.Tx.initialize(vm)
}

// setBaseDatabase sets this block's base database to [db]
func (pb *ProposalBlock) setBaseDatabase(db database.Database) {
	if err := pb.onCommitDB.SetDatabase(db); err != nil {
		pb.vm.Ctx.Log.Error("problem while setting base database: %s", err)
	}
	if err := pb.onAbortDB.SetDatabase(db); err != nil {
		pb.vm.Ctx.Log.Error("problem while setting base database: %s", err)
	}
}

// onCommit should only be called after Verify is called.
// onCommit returns:
//   1. A database that contains the state of the chain assuming this proposal
//      is enacted. (That is, if this block is accepted and followed by an
//      accepted Commit block.)
//   2. A function be be executed when this block's proposal is committed.
//      This function should not write to state.
func (pb *ProposalBlock) onCommit() (*versiondb.Database, func()) {
	return pb.onCommitDB, pb.onCommitFunc
}

// onAbort should only be called after Verify is called.
// onAbort returns a database that contains the state of the chain assuming this
// block's proposal is rejected. (That is, if this block is accepted and
// followed by an accepted Abort block.)
func (pb *ProposalBlock) onAbort() (*versiondb.Database, func()) {
	return pb.onAbortDB, pb.onAbortFunc
}

// Verify this block is valid.
//
// The parent block must either be a Commit or an Abort block.
//
// If this block is valid, this function also sets pas.onCommit and pas.onAbort.
func (pb *ProposalBlock) Verify() error {
	parentIntf := pb.parentBlock()

	// The parent of a proposal block (ie this block) must be a decision block
	parent, ok := parentIntf.(decision)
	if !ok {
		if err := pb.Reject(); err == nil {
			if err := pb.vm.DB.Commit(); err != nil {
				pb.vm.Ctx.Log.Error("error committing Proposal block as rejected: %s", err)
			}
		} else {
			pb.vm.DB.Abort()
		}
		return errInvalidBlockType
	}

	// pdb is the database if this block's parent is accepted
	pdb := parent.onAccept()

	var err TxError
	pb.onCommitDB, pb.onAbortDB, pb.onCommitFunc, pb.onAbortFunc, err = pb.Tx.SemanticVerify(pdb)
	if err != nil {
		// If this block's transaction proposes to advance the timestamp, the transaction may fail
		// verification now but be valid in the future, so don't (permanently) mark the block as rejected.
		if !err.Temporary() {
			if err := pb.Reject(); err == nil {
				if err := pb.vm.DB.Commit(); err != nil {
					pb.vm.Ctx.Log.Error("error committing Proposal block as rejected: %s", err)
				}
			} else {
				pb.vm.DB.Abort()
			}
		}
		return err
	}

	pb.vm.currentBlocks[pb.ID().Key()] = pb
	parentIntf.addChild(pb)
	return nil
}

// Options returns the possible children of this block in preferential order.
func (pb *ProposalBlock) Options() ([2]snowman.Block, error) {
	blockID := pb.ID()

	commit, err := pb.vm.newCommitBlock(blockID)
	if err != nil {
		return [2]snowman.Block{}, err
	}
	abort, err := pb.vm.newAbortBlock(blockID)
	if err != nil {
		return [2]snowman.Block{}, err
	}

	if err := pb.vm.State.PutBlock(pb.vm.DB, commit); err != nil {
		return [2]snowman.Block{}, err
	}
	if err := pb.vm.State.PutBlock(pb.vm.DB, abort); err != nil {
		return [2]snowman.Block{}, err
	}
	if err := pb.vm.DB.Commit(); err != nil {
		return [2]snowman.Block{}, err
	}

	if pb.Tx.InitiallyPrefersCommit() {
		return [2]snowman.Block{commit, abort}, nil
	}
	return [2]snowman.Block{abort, commit}, nil
}

// newProposalBlock creates a new block that proposes to issue a transaction.
// The parent of this block has ID [parentID]. The parent must be a decision block.
// Returns nil if there's an error while creating this block
func (vm *VM) newProposalBlock(parentID ids.ID, tx ProposalTx) (*ProposalBlock, error) {
	pb := &ProposalBlock{
		CommonBlock: CommonBlock{
			Block: core.NewBlock(parentID),
			vm:    vm,
		},
		Tx: tx,
	}

	// We marshal the block in this way (as a Block) so that we can unmarshal
	// it into a Block (rather than a *ProposalBlock)
	block := Block(pb)
	bytes, err := Codec.Marshal(&block)
	if err != nil {
		return nil, err
	}
	pb.Initialize(bytes, vm.SnowmanVM)
	return pb, nil
}
<|endoftext|>


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

package fs

import (
	"bufio"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/user"
	"path/filepath"
	"strings"
)

// readLine returns a single line (without the ending \n)
// from the input buffered reader.
// An error is returned iff there is an error with the
// buffered reader.
func readLine(bufferedReader *bufio.Reader) (string, error) {
	var (
		isPrefix bool  = true
		err      error = nil
		line, ln []byte
	)

	for isPrefix && err == nil {
		line, isPrefix, err = bufferedReader.ReadLine()
		ln = append(ln, line...)
	}

	return string(ln), err
}

// Get all lines of a given file
func GetLines(inFile io.Reader) []string {

	lines := make([]string, 0, 10)
	bufferedReader := bufio.NewReader(inFile)
	line, err := readLine(bufferedReader)
	for err == nil {
		lines = append(lines, line)
		line, err = readLine(bufferedReader)
	}

	return lines
}

func CreateDirectory(path string) bool {
	err := os.MkdirAll(path, 0700)
	return err == nil
}

func Copy(source, target string) (success bool, err error) {

	// check if the source is a file
	if IsFile(source) {
		return CopyFile(source, target)
	}

	// the source must be a directory
	// read the source directory
	sourceEntries, err := ioutil.ReadDir(source)
	if err != nil {
		return false, err
	}

	for _, sourceEntry := range sourceEntries {
		sourceEntryPath := filepath.Join(source, sourceEntry.Name())

		// recurse into the sub-directory
		if sourceEntry.IsDir() {
			nestedTargetPath := filepath.Join(target, sourceEntry.Name())
			if _, err := Copy(sourceEntryPath, nestedTargetPath); err != nil {
				return false, err // abort if an error occurs
			}

			continue
		}

		// copy the file
		targetFilePath := filepath.Join(target, sourceEntry.Name())
		if _, err := CopyFile(sourceEntryPath, targetFilePath); err != nil {
			return false, err // abort if an error occurs
		}
	}

	// if no error occured everything must be ok
	return true, nil
}

func CopyFile(source, target string) (success bool, err error) {
	if !IsFile(source) {
		return false, fmt.Errorf("%q is not a file.", source)
	}

	// open the source file
	sourceReader, readerErr := os.Open(source)
	if readerErr != nil {
		return false, readerErr
	}

	defer sourceReader.Close()

	// prepare the target file
	var targetFileMode os.FileMode = 0600
	if !FileExists(target) {

		// make sure the path to the target file exists
		if _, createFileErr := CreateFile(target); createFileErr != nil {
			return false, fmt.Errorf("Unable to create the target file %q. Error: %s", target, err)
		}
	} else {

		// use the same file mode if the file already exists
		if targetFileInfo, targetFileInfoErr := os.Stat(target); targetFileInfoErr == nil {
			targetFileMode = targetFileInfo.Mode().Perm()
		}

	}

	// open the target file for writing
	targetWriter, writerErr := os.OpenFile(target, os.O_RDWR|os.O_CREATE|os.O_TRUNC, targetFileMode)
	if writerErr != nil {
		return false, writerErr
	}

	defer targetWriter.Close()

	// copy from source to target
	_, copyErr := io.Copy(targetWriter, sourceReader)
	if copyErr != nil {
		return false, err
	}

	return true, nil
}

func CreateFile(filePath string) (success bool, err error) {

	// make sure the parent directory exists
	directory := filepath.Dir(filePath)
	if !DirectoryExists(directory) {
		if !CreateDirectory(directory) {
			return false, fmt.Errorf("Cannot create the directory for the given file %q.", filePath)
		}
	}

	// create the file
	if _, err := os.Create(filePath); err != nil {
		return false, fmt.Errorf("Could not create file %q. Error: ", filePath, err)
	}

	return true, nil
}

func PathExists(path string) bool {
	if strings.TrimSpace(path) == "" {
		return false
	}

	if _, err := os.Stat(path); err != nil && os.IsNotExist(err) {
		return false
	}

	return true
}

func FileExists(path string) bool {
	if strings.TrimSpace(path) == "" {
		return false
	}

	file, err := os.Stat(path)
	if err != nil && os.IsNotExist(err) {
		return false
	}

	return !file.IsDir()
}

func DirectoryExists(path string) bool {
	if strings.TrimSpace(path) == "" {
		return false
	}

	file, err := os.Stat(path)
	if err != nil && os.IsNotExist(err) {
		return false
	}

	return file.IsDir()
}

func IsFile(path string) bool {

	fileInfo, err := os.Stat(path)
	if err != nil {
		return false
	}

	return fileInfo.IsDir() == false
}

func IsDirectory(path string) bool {

	fileInfo, err := os.Stat(path)
	if err != nil {
		return false
	}

	return fileInfo.IsDir()
}

func GetUserHomeDirectory() (string, error) {

	usr, err := user.Current()
	if err != nil {
		return "", err
	}

	return filepath.Clean(usr.HomeDir), nil
}

func GetAllFilesInDirectory(path string) []string {

	files := make([]string, 0)

	if !IsDirectory(path) {
		return files
	}

	directoryEntries, err := ioutil.ReadDir(path)
	if err != nil {
		return files
	}

	for _, entry := range directoryEntries {

		entryPath := filepath.Join(path, entry.Name())

		// recurse
		if entry.IsDir() {
			files = append(files, GetAllFilesInDirectory(entryPath)...)
		}

		files = append(files, entryPath)
	}

	return files
}
<commit_msg>Bug fix for the new GetAllFilesInDirectory function. The recursion clause did not end the current loop cycle and added the directory path to the files list which causes prolems with the backup action.<commit_after>// Copyright 2013 Andreas Koch. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package fs

import (
	"bufio"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/user"
	"path/filepath"
	"strings"
)

// readLine returns a single line (without the ending \n)
// from the input buffered reader.
// An error is returned iff there is an error with the
// buffered reader.
func readLine(bufferedReader *bufio.Reader) (string, error) {
	var (
		isPrefix bool  = true
		err      error = nil
		line, ln []byte
	)

	for isPrefix && err == nil {
		line, isPrefix, err = bufferedReader.ReadLine()
		ln = append(ln, line...)
	}

	return string(ln), err
}

// Get all lines of a given file
func GetLines(inFile io.Reader) []string {

	lines := make([]string, 0, 10)
	bufferedReader := bufio.NewReader(inFile)
	line, err := readLine(bufferedReader)
	for err == nil {
		lines = append(lines, line)
		line, err = readLine(bufferedReader)
	}

	return lines
}

func CreateDirectory(path string) bool {
	err := os.MkdirAll(path, 0700)
	return err == nil
}

func Copy(source, target string) (success bool, err error) {

	// check if the source is a file
	if IsFile(source) {
		return CopyFile(source, target)
	}

	// the source must be a directory
	// read the source directory
	sourceEntries, err := ioutil.ReadDir(source)
	if err != nil {
		return false, err
	}

	for _, sourceEntry := range sourceEntries {
		sourceEntryPath := filepath.Join(source, sourceEntry.Name())

		// recurse into the sub-directory
		if sourceEntry.IsDir() {
			nestedTargetPath := filepath.Join(target, sourceEntry.Name())
			if _, err := Copy(sourceEntryPath, nestedTargetPath); err != nil {
				return false, err // abort if an error occurs
			}

			continue
		}

		// copy the file
		targetFilePath := filepath.Join(target, sourceEntry.Name())
		if _, err := CopyFile(sourceEntryPath, targetFilePath); err != nil {
			return false, err // abort if an error occurs
		}
	}

	// if no error occured everything must be ok
	return true, nil
}

func CopyFile(source, target string) (success bool, err error) {
	if !IsFile(source) {
		return false, fmt.Errorf("%q is not a file.", source)
	}

	// open the source file
	sourceReader, readerErr := os.Open(source)
	if readerErr != nil {
		return false, readerErr
	}

	defer sourceReader.Close()

	// prepare the target file
	var targetFileMode os.FileMode = 0600
	if !FileExists(target) {

		// make sure the path to the target file exists
		if _, createFileErr := CreateFile(target); createFileErr != nil {
			return false, fmt.Errorf("Unable to create the target file %q. Error: %s", target, err)
		}
	} else {

		// use the same file mode if the file already exists
		if targetFileInfo, targetFileInfoErr := os.Stat(target); targetFileInfoErr == nil {
			targetFileMode = targetFileInfo.Mode().Perm()
		}

	}

	// open the target file for writing
	targetWriter, writerErr := os.OpenFile(target, os.O_RDWR|os.O_CREATE|os.O_TRUNC, targetFileMode)
	if writerErr != nil {
		return false, writerErr
	}

	defer targetWriter.Close()

	// copy from source to target
	_, copyErr := io.Copy(targetWriter, sourceReader)
	if copyErr != nil {
		return false, err
	}

	return true, nil
}

func CreateFile(filePath string) (success bool, err error) {

	// make sure the parent directory exists
	directory := filepath.Dir(filePath)
	if !DirectoryExists(directory) {
		if !CreateDirectory(directory) {
			return false, fmt.Errorf("Cannot create the directory for the given file %q.", filePath)
		}
	}

	// create the file
	if _, err := os.Create(filePath); err != nil {
		return false, fmt.Errorf("Could not create file %q. Error: ", filePath, err)
	}

	return true, nil
}

func PathExists(path string) bool {
	if strings.TrimSpace(path) == "" {
		return false
	}

	if _, err := os.Stat(path); err != nil && os.IsNotExist(err) {
		return false
	}

	return true
}

func FileExists(path string) bool {
	if strings.TrimSpace(path) == "" {
		return false
	}

	file, err := os.Stat(path)
	if err != nil && os.IsNotExist(err) {
		return false
	}

	return !file.IsDir()
}

func DirectoryExists(path string) bool {
	if strings.TrimSpace(path) == "" {
		return false
	}

	file, err := os.Stat(path)
	if err != nil && os.IsNotExist(err) {
		return false
	}

	return file.IsDir()
}

func IsFile(path string) bool {

	fileInfo, err := os.Stat(path)
	if err != nil {
		return false
	}

	return fileInfo.IsDir() == false
}

func IsDirectory(path string) bool {

	fileInfo, err := os.Stat(path)
	if err != nil {
		return false
	}

	return fileInfo.IsDir()
}

func GetUserHomeDirectory() (string, error) {

	usr, err := user.Current()
	if err != nil {
		return "", err
	}

	return filepath.Clean(usr.HomeDir), nil
}

func GetAllFilesInDirectory(path string) []string {

	files := make([]string, 0)

	if !IsDirectory(path) {
		return files
	}

	directoryEntries, err := ioutil.ReadDir(path)
	if err != nil {
		return files
	}

	for _, entry := range directoryEntries {

		entryPath := filepath.Join(path, entry.Name())

		// recurse
		if entry.IsDir() {
			files = append(files, GetAllFilesInDirectory(entryPath)...)
			continue
		}

		files = append(files, entryPath)
	}

	return files
}
<|endoftext|>


<commit_before>package httputil

import (
	"context"
	"net"
	"net/http"
	"time"
)

func NewPooledUnixTransport(sockFile string) *http.Transport {
	return &http.Transport{
		DialContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
			d := net.Dialer{
				Timeout:   30 * time.Second,
				KeepAlive: 30 * time.Second,
				DualStack: true,
			}
			return d.DialContext(ctx, "unix", sockFile)
		},
		MaxIdleConns:          100,
		IdleConnTimeout:       90 * time.Second,
		TLSHandshakeTimeout:   10 * time.Second,
		ExpectContinueTimeout: 1 * time.Second,
	}
}

func ListenAndServeUnix(srv *http.Server, addr string) error {
	ln, err := net.Listen("unix", addr)
	if err != nil {
		return err
	}
	// TODO: do we need tcpKeepAliveListener like the default tcp ListenAndServe?
	return srv.Serve(ln)
}
<commit_msg>[httputil] Add unix transport without pool<commit_after>package httputil

import (
	"context"
	"net"
	"net/http"
	"time"
)

// NewUnPooledUnixTransport uses NewPooledUnixTransport with keep alive and idle connection disabled
func NewUnPooledUnixTransport(socketFile string) *http.Transport {
	tr := NewPooledUnixTransport(socketFile)
	tr.DisableKeepAlives = true
	tr.MaxIdleConnsPerHost = -1
	return tr
}

// NewPooledUnixTransport creates transport for unix domain socket i.e. /var/run/docker.sock
func NewPooledUnixTransport(sockFile string) *http.Transport {
	return &http.Transport{
		DialContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
			d := net.Dialer{
				Timeout:   30 * time.Second,
				KeepAlive: 30 * time.Second,
				DualStack: true,
			}
			return d.DialContext(ctx, "unix", sockFile)
		},
		MaxIdleConns:          100,
		IdleConnTimeout:       90 * time.Second,
		TLSHandshakeTimeout:   10 * time.Second,
		ExpectContinueTimeout: 1 * time.Second,
	}
}

func ListenAndServeUnix(srv *http.Server, addr string) error {
	ln, err := net.Listen("unix", addr)
	if err != nil {
		return err
	}
	// TODO: do we need tcpKeepAliveListener like the default tcp ListenAndServe?
	return srv.Serve(ln)
}
<|endoftext|>


<commit_before>package aws

import (
	"fmt"
	"log"
	"strconv"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/arn"
	"github.com/aws/aws-sdk-go/service/apigateway"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/helper/schema"
	"github.com/hashicorp/terraform/helper/structure"
	"github.com/hashicorp/terraform/helper/validation"
)

func resourceAwsApiGatewayRestApi() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsApiGatewayRestApiCreate,
		Read:   resourceAwsApiGatewayRestApiRead,
		Update: resourceAwsApiGatewayRestApiUpdate,
		Delete: resourceAwsApiGatewayRestApiDelete,
		Importer: &schema.ResourceImporter{
			State: schema.ImportStatePassthrough,
		},

		Schema: map[string]*schema.Schema{
			"name": {
				Type:     schema.TypeString,
				Required: true,
			},

			"description": {
				Type:     schema.TypeString,
				Optional: true,
			},

			"api_key_source": {
				Type:     schema.TypeString,
				Optional: true,
				Default:  "HEADER",
			},

			"policy": {
				Type:             schema.TypeString,
				Optional:         true,
				ValidateFunc:     validation.ValidateJsonString,
				DiffSuppressFunc: suppressEquivalentAwsPolicyDiffs,
			},

			"binary_media_types": {
				Type:     schema.TypeList,
				Optional: true,
				Elem:     &schema.Schema{Type: schema.TypeString},
			},

			"body": {
				Type:     schema.TypeString,
				Optional: true,
			},

			"minimum_compression_size": {
				Type:         schema.TypeInt,
				Optional:     true,
				Default:      -1,
				ValidateFunc: validation.IntBetween(-1, 10485760),
			},

			"root_resource_id": {
				Type:     schema.TypeString,
				Computed: true,
			},

			"created_date": {
				Type:     schema.TypeString,
				Computed: true,
			},
			"execution_arn": {
				Type:     schema.TypeString,
				Computed: true,
			},

			"endpoint_configuration": {
				Type:     schema.TypeList,
				Optional: true,
				Computed: true,
				MinItems: 1,
				MaxItems: 1,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"types": {
							Type:     schema.TypeList,
							Required: true,
							MinItems: 1,
							MaxItems: 1,
							Elem: &schema.Schema{
								Type: schema.TypeString,
								ValidateFunc: validation.StringInSlice([]string{
									apigateway.EndpointTypeEdge,
									apigateway.EndpointTypeRegional,
									apigateway.EndpointTypePrivate,
								}, false),
							},
						},
					},
				},
			},
		},
	}
}

func resourceAwsApiGatewayRestApiCreate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).apigateway
	log.Printf("[DEBUG] Creating API Gateway")

	var description *string
	if d.Get("description").(string) != "" {
		description = aws.String(d.Get("description").(string))
	}

	params := &apigateway.CreateRestApiInput{
		Name:        aws.String(d.Get("name").(string)),
		Description: description,
	}

	if v, ok := d.GetOk("endpoint_configuration"); ok {
		params.EndpointConfiguration = expandApiGatewayEndpointConfiguration(v.([]interface{}))
	}

	if v, ok := d.GetOk("api_key_source"); ok && v.(string) != "" {
		params.ApiKeySource = aws.String(v.(string))
	}

	if v, ok := d.GetOk("policy"); ok && v.(string) != "" {
		params.Policy = aws.String(v.(string))
	}

	binaryMediaTypes, binaryMediaTypesOk := d.GetOk("binary_media_types")
	if binaryMediaTypesOk {
		params.BinaryMediaTypes = expandStringList(binaryMediaTypes.([]interface{}))
	}

	minimumCompressionSize := d.Get("minimum_compression_size").(int)
	if minimumCompressionSize > -1 {
		params.MinimumCompressionSize = aws.Int64(int64(minimumCompressionSize))
	}

	gateway, err := conn.CreateRestApi(params)
	if err != nil {
		return fmt.Errorf("Error creating API Gateway: %s", err)
	}

	d.SetId(*gateway.Id)

	if body, ok := d.GetOk("body"); ok {
		log.Printf("[DEBUG] Initializing API Gateway from OpenAPI spec %s", d.Id())
		_, err := conn.PutRestApi(&apigateway.PutRestApiInput{
			RestApiId: gateway.Id,
			Mode:      aws.String(apigateway.PutModeOverwrite),
			Body:      []byte(body.(string)),
		})
		if err != nil {
			return fmt.Errorf("error creating API Gateway specification: %s", err)
		}
	}

	return resourceAwsApiGatewayRestApiRead(d, meta)
}

func resourceAwsApiGatewayRestApiRead(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).apigateway
	log.Printf("[DEBUG] Reading API Gateway %s", d.Id())

	api, err := conn.GetRestApi(&apigateway.GetRestApiInput{
		RestApiId: aws.String(d.Id()),
	})
	if isAWSErr(err, apigateway.ErrCodeNotFoundException, "") {
		log.Printf("[WARN] API Gateway (%s) not found, removing from state", d.Id())
		d.SetId("")
		return nil
	}
	if err != nil {
		return fmt.Errorf("error reading API Gateway REST API (%s): %s", d.Id(), err)
	}

	getResourcesInput := &apigateway.GetResourcesInput{
		RestApiId: aws.String(d.Id()),
	}
	err = conn.GetResourcesPages(getResourcesInput, func(page *apigateway.GetResourcesOutput, lastPage bool) bool {
		for _, item := range page.Items {
			if aws.StringValue(item.Path) == "/" {
				d.Set("root_resource_id", item.Id)
				return false
			}
		}
		return !lastPage
	})
	if err != nil {
		return fmt.Errorf("error reading API Gateway REST API (%s) resources: %s", d.Id(), err)
	}

	d.Set("name", api.Name)
	d.Set("description", api.Description)
	d.Set("api_key_source", api.ApiKeySource)

	// The API returns policy as an escaped JSON string
	// {\\\"Version\\\":\\\"2012-10-17\\\",...}
	// The string must be normalized before unquoting as it may contain escaped
	// forward slashes in CIDR blocks, which will break strconv.Unquote

	// I'm not sure why it needs to be wrapped with double quotes first, but it does
	normalized_policy, err := structure.NormalizeJsonString(`"` + aws.StringValue(api.Policy) + `"`)
	if err != nil {
		fmt.Printf("error normalizing policy JSON: %s\n", err)
	}
	policy, err := strconv.Unquote(normalized_policy)
	if err != nil {
		return fmt.Errorf("error unescaping policy: %s", err)
	}
	d.Set("policy", policy)

	d.Set("binary_media_types", api.BinaryMediaTypes)

	arn := arn.ARN{
		Partition: meta.(*AWSClient).partition,
		Service:   "execute-api",
		Region:    meta.(*AWSClient).region,
		AccountID: meta.(*AWSClient).accountid,
		Resource:  d.Id(),
	}.String()
	d.Set("execution_arn", arn)

	if api.MinimumCompressionSize == nil {
		d.Set("minimum_compression_size", -1)
	} else {
		d.Set("minimum_compression_size", api.MinimumCompressionSize)
	}
	if err := d.Set("created_date", api.CreatedDate.Format(time.RFC3339)); err != nil {
		log.Printf("[DEBUG] Error setting created_date: %s", err)
	}

	if err := d.Set("endpoint_configuration", flattenApiGatewayEndpointConfiguration(api.EndpointConfiguration)); err != nil {
		return fmt.Errorf("error setting endpoint_configuration: %s", err)
	}

	return nil
}

func resourceAwsApiGatewayRestApiUpdateOperations(d *schema.ResourceData) []*apigateway.PatchOperation {
	operations := make([]*apigateway.PatchOperation, 0)

	if d.HasChange("name") {
		operations = append(operations, &apigateway.PatchOperation{
			Op:    aws.String("replace"),
			Path:  aws.String("/name"),
			Value: aws.String(d.Get("name").(string)),
		})
	}

	if d.HasChange("description") {
		operations = append(operations, &apigateway.PatchOperation{
			Op:    aws.String("replace"),
			Path:  aws.String("/description"),
			Value: aws.String(d.Get("description").(string)),
		})
	}

	if d.HasChange("api_key_source") {
		operations = append(operations, &apigateway.PatchOperation{
			Op:    aws.String("replace"),
			Path:  aws.String("/apiKeySource"),
			Value: aws.String(d.Get("api_key_source").(string)),
		})
	}

	if d.HasChange("policy") {
		operations = append(operations, &apigateway.PatchOperation{
			Op:    aws.String("replace"),
			Path:  aws.String("/policy"),
			Value: aws.String(d.Get("policy").(string)),
		})
	}

	if d.HasChange("minimum_compression_size") {
		minimumCompressionSize := d.Get("minimum_compression_size").(int)
		var value string
		if minimumCompressionSize > -1 {
			value = strconv.Itoa(minimumCompressionSize)
		}
		operations = append(operations, &apigateway.PatchOperation{
			Op:    aws.String("replace"),
			Path:  aws.String("/minimumCompressionSize"),
			Value: aws.String(value),
		})
	}

	if d.HasChange("binary_media_types") {
		o, n := d.GetChange("binary_media_types")
		prefix := "binaryMediaTypes"

		old := o.([]interface{})
		new := n.([]interface{})

		// Remove every binary media types. Simpler to remove and add new ones,
		// since there are no replacings.
		for _, v := range old {
			operations = append(operations, &apigateway.PatchOperation{
				Op:   aws.String("remove"),
				Path: aws.String(fmt.Sprintf("/%s/%s", prefix, escapeJsonPointer(v.(string)))),
			})
		}

		// Handle additions
		if len(new) > 0 {
			for _, v := range new {
				operations = append(operations, &apigateway.PatchOperation{
					Op:   aws.String("add"),
					Path: aws.String(fmt.Sprintf("/%s/%s", prefix, escapeJsonPointer(v.(string)))),
				})
			}
		}
	}

	if d.HasChange("endpoint_configuration.0.types") {
		// The REST API must have an endpoint type.
		// If attempting to remove the configuration, do nothing.
		if v, ok := d.GetOk("endpoint_configuration"); ok && len(v.([]interface{})) > 0 {
			m := v.([]interface{})[0].(map[string]interface{})

			operations = append(operations, &apigateway.PatchOperation{
				Op:    aws.String("replace"),
				Path:  aws.String("/endpointConfiguration/types/0"),
				Value: aws.String(m["types"].([]interface{})[0].(string)),
			})
		}
	}

	return operations
}

func resourceAwsApiGatewayRestApiUpdate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).apigateway
	log.Printf("[DEBUG] Updating API Gateway %s", d.Id())

	if d.HasChange("body") {
		if body, ok := d.GetOk("body"); ok {
			log.Printf("[DEBUG] Updating API Gateway from OpenAPI spec: %s", d.Id())
			_, err := conn.PutRestApi(&apigateway.PutRestApiInput{
				RestApiId: aws.String(d.Id()),
				Mode:      aws.String(apigateway.PutModeOverwrite),
				Body:      []byte(body.(string)),
			})
			if err != nil {
				return fmt.Errorf("error updating API Gateway specification: %s", err)
			}
		}
	}

	_, err := conn.UpdateRestApi(&apigateway.UpdateRestApiInput{
		RestApiId:       aws.String(d.Id()),
		PatchOperations: resourceAwsApiGatewayRestApiUpdateOperations(d),
	})

	if err != nil {
		return err
	}
	log.Printf("[DEBUG] Updated API Gateway %s", d.Id())

	return resourceAwsApiGatewayRestApiRead(d, meta)
}

func resourceAwsApiGatewayRestApiDelete(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).apigateway
	log.Printf("[DEBUG] Deleting API Gateway: %s", d.Id())

	return resource.Retry(10*time.Minute, func() *resource.RetryError {
		_, err := conn.DeleteRestApi(&apigateway.DeleteRestApiInput{
			RestApiId: aws.String(d.Id()),
		})
		if err == nil {
			return nil
		}

		if isAWSErr(err, apigateway.ErrCodeNotFoundException, "") {
			return nil
		}

		return resource.NonRetryableError(err)
	})
}

func expandApiGatewayEndpointConfiguration(l []interface{}) *apigateway.EndpointConfiguration {
	if len(l) == 0 {
		return nil
	}

	m := l[0].(map[string]interface{})

	endpointConfiguration := &apigateway.EndpointConfiguration{
		Types: expandStringList(m["types"].([]interface{})),
	}

	return endpointConfiguration
}

func flattenApiGatewayEndpointConfiguration(endpointConfiguration *apigateway.EndpointConfiguration) []interface{} {
	if endpointConfiguration == nil {
		return []interface{}{}
	}

	m := map[string]interface{}{
		"types": flattenStringList(endpointConfiguration.Types),
	}

	return []interface{}{m}
}
<commit_msg>resource/aws_api_gateway_rest_api: Remove extraneous timeout for resource deletion<commit_after>package aws

import (
	"fmt"
	"log"
	"strconv"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/arn"
	"github.com/aws/aws-sdk-go/service/apigateway"
	"github.com/hashicorp/terraform/helper/schema"
	"github.com/hashicorp/terraform/helper/structure"
	"github.com/hashicorp/terraform/helper/validation"
)

func resourceAwsApiGatewayRestApi() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsApiGatewayRestApiCreate,
		Read:   resourceAwsApiGatewayRestApiRead,
		Update: resourceAwsApiGatewayRestApiUpdate,
		Delete: resourceAwsApiGatewayRestApiDelete,
		Importer: &schema.ResourceImporter{
			State: schema.ImportStatePassthrough,
		},

		Schema: map[string]*schema.Schema{
			"name": {
				Type:     schema.TypeString,
				Required: true,
			},

			"description": {
				Type:     schema.TypeString,
				Optional: true,
			},

			"api_key_source": {
				Type:     schema.TypeString,
				Optional: true,
				Default:  "HEADER",
			},

			"policy": {
				Type:             schema.TypeString,
				Optional:         true,
				ValidateFunc:     validation.ValidateJsonString,
				DiffSuppressFunc: suppressEquivalentAwsPolicyDiffs,
			},

			"binary_media_types": {
				Type:     schema.TypeList,
				Optional: true,
				Elem:     &schema.Schema{Type: schema.TypeString},
			},

			"body": {
				Type:     schema.TypeString,
				Optional: true,
			},

			"minimum_compression_size": {
				Type:         schema.TypeInt,
				Optional:     true,
				Default:      -1,
				ValidateFunc: validation.IntBetween(-1, 10485760),
			},

			"root_resource_id": {
				Type:     schema.TypeString,
				Computed: true,
			},

			"created_date": {
				Type:     schema.TypeString,
				Computed: true,
			},
			"execution_arn": {
				Type:     schema.TypeString,
				Computed: true,
			},

			"endpoint_configuration": {
				Type:     schema.TypeList,
				Optional: true,
				Computed: true,
				MinItems: 1,
				MaxItems: 1,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"types": {
							Type:     schema.TypeList,
							Required: true,
							MinItems: 1,
							MaxItems: 1,
							Elem: &schema.Schema{
								Type: schema.TypeString,
								ValidateFunc: validation.StringInSlice([]string{
									apigateway.EndpointTypeEdge,
									apigateway.EndpointTypeRegional,
									apigateway.EndpointTypePrivate,
								}, false),
							},
						},
					},
				},
			},
		},
	}
}

func resourceAwsApiGatewayRestApiCreate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).apigateway
	log.Printf("[DEBUG] Creating API Gateway")

	var description *string
	if d.Get("description").(string) != "" {
		description = aws.String(d.Get("description").(string))
	}

	params := &apigateway.CreateRestApiInput{
		Name:        aws.String(d.Get("name").(string)),
		Description: description,
	}

	if v, ok := d.GetOk("endpoint_configuration"); ok {
		params.EndpointConfiguration = expandApiGatewayEndpointConfiguration(v.([]interface{}))
	}

	if v, ok := d.GetOk("api_key_source"); ok && v.(string) != "" {
		params.ApiKeySource = aws.String(v.(string))
	}

	if v, ok := d.GetOk("policy"); ok && v.(string) != "" {
		params.Policy = aws.String(v.(string))
	}

	binaryMediaTypes, binaryMediaTypesOk := d.GetOk("binary_media_types")
	if binaryMediaTypesOk {
		params.BinaryMediaTypes = expandStringList(binaryMediaTypes.([]interface{}))
	}

	minimumCompressionSize := d.Get("minimum_compression_size").(int)
	if minimumCompressionSize > -1 {
		params.MinimumCompressionSize = aws.Int64(int64(minimumCompressionSize))
	}

	gateway, err := conn.CreateRestApi(params)
	if err != nil {
		return fmt.Errorf("Error creating API Gateway: %s", err)
	}

	d.SetId(*gateway.Id)

	if body, ok := d.GetOk("body"); ok {
		log.Printf("[DEBUG] Initializing API Gateway from OpenAPI spec %s", d.Id())
		_, err := conn.PutRestApi(&apigateway.PutRestApiInput{
			RestApiId: gateway.Id,
			Mode:      aws.String(apigateway.PutModeOverwrite),
			Body:      []byte(body.(string)),
		})
		if err != nil {
			return fmt.Errorf("error creating API Gateway specification: %s", err)
		}
	}

	return resourceAwsApiGatewayRestApiRead(d, meta)
}

func resourceAwsApiGatewayRestApiRead(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).apigateway
	log.Printf("[DEBUG] Reading API Gateway %s", d.Id())

	api, err := conn.GetRestApi(&apigateway.GetRestApiInput{
		RestApiId: aws.String(d.Id()),
	})
	if isAWSErr(err, apigateway.ErrCodeNotFoundException, "") {
		log.Printf("[WARN] API Gateway (%s) not found, removing from state", d.Id())
		d.SetId("")
		return nil
	}
	if err != nil {
		return fmt.Errorf("error reading API Gateway REST API (%s): %s", d.Id(), err)
	}

	getResourcesInput := &apigateway.GetResourcesInput{
		RestApiId: aws.String(d.Id()),
	}
	err = conn.GetResourcesPages(getResourcesInput, func(page *apigateway.GetResourcesOutput, lastPage bool) bool {
		for _, item := range page.Items {
			if aws.StringValue(item.Path) == "/" {
				d.Set("root_resource_id", item.Id)
				return false
			}
		}
		return !lastPage
	})
	if err != nil {
		return fmt.Errorf("error reading API Gateway REST API (%s) resources: %s", d.Id(), err)
	}

	d.Set("name", api.Name)
	d.Set("description", api.Description)
	d.Set("api_key_source", api.ApiKeySource)

	// The API returns policy as an escaped JSON string
	// {\\\"Version\\\":\\\"2012-10-17\\\",...}
	// The string must be normalized before unquoting as it may contain escaped
	// forward slashes in CIDR blocks, which will break strconv.Unquote

	// I'm not sure why it needs to be wrapped with double quotes first, but it does
	normalized_policy, err := structure.NormalizeJsonString(`"` + aws.StringValue(api.Policy) + `"`)
	if err != nil {
		fmt.Printf("error normalizing policy JSON: %s\n", err)
	}
	policy, err := strconv.Unquote(normalized_policy)
	if err != nil {
		return fmt.Errorf("error unescaping policy: %s", err)
	}
	d.Set("policy", policy)

	d.Set("binary_media_types", api.BinaryMediaTypes)

	arn := arn.ARN{
		Partition: meta.(*AWSClient).partition,
		Service:   "execute-api",
		Region:    meta.(*AWSClient).region,
		AccountID: meta.(*AWSClient).accountid,
		Resource:  d.Id(),
	}.String()
	d.Set("execution_arn", arn)

	if api.MinimumCompressionSize == nil {
		d.Set("minimum_compression_size", -1)
	} else {
		d.Set("minimum_compression_size", api.MinimumCompressionSize)
	}
	if err := d.Set("created_date", api.CreatedDate.Format(time.RFC3339)); err != nil {
		log.Printf("[DEBUG] Error setting created_date: %s", err)
	}

	if err := d.Set("endpoint_configuration", flattenApiGatewayEndpointConfiguration(api.EndpointConfiguration)); err != nil {
		return fmt.Errorf("error setting endpoint_configuration: %s", err)
	}

	return nil
}

func resourceAwsApiGatewayRestApiUpdateOperations(d *schema.ResourceData) []*apigateway.PatchOperation {
	operations := make([]*apigateway.PatchOperation, 0)

	if d.HasChange("name") {
		operations = append(operations, &apigateway.PatchOperation{
			Op:    aws.String("replace"),
			Path:  aws.String("/name"),
			Value: aws.String(d.Get("name").(string)),
		})
	}

	if d.HasChange("description") {
		operations = append(operations, &apigateway.PatchOperation{
			Op:    aws.String("replace"),
			Path:  aws.String("/description"),
			Value: aws.String(d.Get("description").(string)),
		})
	}

	if d.HasChange("api_key_source") {
		operations = append(operations, &apigateway.PatchOperation{
			Op:    aws.String("replace"),
			Path:  aws.String("/apiKeySource"),
			Value: aws.String(d.Get("api_key_source").(string)),
		})
	}

	if d.HasChange("policy") {
		operations = append(operations, &apigateway.PatchOperation{
			Op:    aws.String("replace"),
			Path:  aws.String("/policy"),
			Value: aws.String(d.Get("policy").(string)),
		})
	}

	if d.HasChange("minimum_compression_size") {
		minimumCompressionSize := d.Get("minimum_compression_size").(int)
		var value string
		if minimumCompressionSize > -1 {
			value = strconv.Itoa(minimumCompressionSize)
		}
		operations = append(operations, &apigateway.PatchOperation{
			Op:    aws.String("replace"),
			Path:  aws.String("/minimumCompressionSize"),
			Value: aws.String(value),
		})
	}

	if d.HasChange("binary_media_types") {
		o, n := d.GetChange("binary_media_types")
		prefix := "binaryMediaTypes"

		old := o.([]interface{})
		new := n.([]interface{})

		// Remove every binary media types. Simpler to remove and add new ones,
		// since there are no replacings.
		for _, v := range old {
			operations = append(operations, &apigateway.PatchOperation{
				Op:   aws.String("remove"),
				Path: aws.String(fmt.Sprintf("/%s/%s", prefix, escapeJsonPointer(v.(string)))),
			})
		}

		// Handle additions
		if len(new) > 0 {
			for _, v := range new {
				operations = append(operations, &apigateway.PatchOperation{
					Op:   aws.String("add"),
					Path: aws.String(fmt.Sprintf("/%s/%s", prefix, escapeJsonPointer(v.(string)))),
				})
			}
		}
	}

	if d.HasChange("endpoint_configuration.0.types") {
		// The REST API must have an endpoint type.
		// If attempting to remove the configuration, do nothing.
		if v, ok := d.GetOk("endpoint_configuration"); ok && len(v.([]interface{})) > 0 {
			m := v.([]interface{})[0].(map[string]interface{})

			operations = append(operations, &apigateway.PatchOperation{
				Op:    aws.String("replace"),
				Path:  aws.String("/endpointConfiguration/types/0"),
				Value: aws.String(m["types"].([]interface{})[0].(string)),
			})
		}
	}

	return operations
}

func resourceAwsApiGatewayRestApiUpdate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).apigateway
	log.Printf("[DEBUG] Updating API Gateway %s", d.Id())

	if d.HasChange("body") {
		if body, ok := d.GetOk("body"); ok {
			log.Printf("[DEBUG] Updating API Gateway from OpenAPI spec: %s", d.Id())
			_, err := conn.PutRestApi(&apigateway.PutRestApiInput{
				RestApiId: aws.String(d.Id()),
				Mode:      aws.String(apigateway.PutModeOverwrite),
				Body:      []byte(body.(string)),
			})
			if err != nil {
				return fmt.Errorf("error updating API Gateway specification: %s", err)
			}
		}
	}

	_, err := conn.UpdateRestApi(&apigateway.UpdateRestApiInput{
		RestApiId:       aws.String(d.Id()),
		PatchOperations: resourceAwsApiGatewayRestApiUpdateOperations(d),
	})

	if err != nil {
		return err
	}
	log.Printf("[DEBUG] Updated API Gateway %s", d.Id())

	return resourceAwsApiGatewayRestApiRead(d, meta)
}

func resourceAwsApiGatewayRestApiDelete(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).apigateway

	input := &apigateway.DeleteRestApiInput{
		RestApiId: aws.String(d.Id()),
	}

	log.Printf("[DEBUG] Deleting API Gateway: %s", input)
	_, err := conn.DeleteRestApi(input)

	if isAWSErr(err, apigateway.ErrCodeNotFoundException, "") {
		return nil
	}

	if err != nil {
		return fmt.Errorf("error deleting API Gateway (%s): %s", d.Id(), err)
	}

	return nil
}

func expandApiGatewayEndpointConfiguration(l []interface{}) *apigateway.EndpointConfiguration {
	if len(l) == 0 {
		return nil
	}

	m := l[0].(map[string]interface{})

	endpointConfiguration := &apigateway.EndpointConfiguration{
		Types: expandStringList(m["types"].([]interface{})),
	}

	return endpointConfiguration
}

func flattenApiGatewayEndpointConfiguration(endpointConfiguration *apigateway.EndpointConfiguration) []interface{} {
	if endpointConfiguration == nil {
		return []interface{}{}
	}

	m := map[string]interface{}{
		"types": flattenStringList(endpointConfiguration.Types),
	}

	return []interface{}{m}
}
<|endoftext|>


<commit_before>package aws

import (
	"fmt"

	"encoding/json"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/batch"
	"github.com/hashicorp/terraform/helper/schema"
	"github.com/hashicorp/terraform/helper/structure"
	"github.com/hashicorp/terraform/helper/validation"
)

func resourceAwsBatchJobDefinition() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsBatchJobDefinitionCreate,
		Read:   resourceAwsBatchJobDefinitionRead,
		Delete: resourceAwsBatchJobDefinitionDelete,

		Schema: map[string]*schema.Schema{
			"name": {
				Type:         schema.TypeString,
				Required:     true,
				ForceNew:     true,
				ValidateFunc: validateBatchName,
			},
			"container_properties": {
				Type:     schema.TypeString,
				Optional: true,
				ForceNew: true,
				StateFunc: func(v interface{}) string {
					json, _ := structure.NormalizeJsonString(v)
					return json
				},
				DiffSuppressFunc: suppressEquivalentJsonDiffs,
				ValidateFunc:     validateAwsBatchJobContainerProperties,
			},
			"parameters": {
				Type:     schema.TypeMap,
				Optional: true,
				ForceNew: true,
				Elem:     schema.TypeString,
			},
			"retry_strategy": {
				Type:     schema.TypeList,
				Optional: true,
				ForceNew: true,
				MaxItems: 1,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"attempts": {
							Type:         schema.TypeInt,
							Optional:     true,
							ValidateFunc: validation.IntBetween(1, 10),
						},
					},
				},
			},
			"timeout": {
				Type:     schema.TypeList,
				Optional: true,
				ForceNew: true,
				MaxItems: 1,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"attempt_duration_seconds": {
							Type:         schema.TypeInt,
							Optional:     true,
							ForceNew:     true,
							ValidateFunc: validation.IntAtLeast(60),
						},
					},
				},
			},
			"type": {
				Type:         schema.TypeString,
				Required:     true,
				ForceNew:     true,
				ValidateFunc: validation.StringInSlice([]string{batch.JobDefinitionTypeContainer}, true),
			},
			"revision": {
				Type:     schema.TypeInt,
				Computed: true,
			},
			"arn": {
				Type:     schema.TypeString,
				Computed: true,
			},
		},
	}
}

func resourceAwsBatchJobDefinitionCreate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).batchconn
	name := d.Get("name").(string)

	input := &batch.RegisterJobDefinitionInput{
		JobDefinitionName: aws.String(name),
		Type:              aws.String(d.Get("type").(string)),
	}

	if v, ok := d.GetOk("container_properties"); ok {
		props, err := expandBatchJobContainerProperties(v.(string))
		if err != nil {
			return fmt.Errorf("%s %q", err, name)
		}
		input.ContainerProperties = props
	}

	if v, ok := d.GetOk("parameters"); ok {
		input.Parameters = expandJobDefinitionParameters(v.(map[string]interface{}))
	}

	if v, ok := d.GetOk("retry_strategy"); ok {
		input.RetryStrategy = expandJobDefinitionRetryStrategy(v.([]interface{}))
	}

	if v, ok := d.GetOk("timeout"); ok {
		input.Timeout = expandJobDefinitionTimeout(v.([]interface{}))
	}

	out, err := conn.RegisterJobDefinition(input)
	if err != nil {
		return fmt.Errorf("%s %q", err, name)
	}
	d.SetId(*out.JobDefinitionArn)
	d.Set("arn", out.JobDefinitionArn)
	return resourceAwsBatchJobDefinitionRead(d, meta)
}

func resourceAwsBatchJobDefinitionRead(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).batchconn
	arn := d.Get("arn").(string)
	job, err := getJobDefinition(conn, arn)
	if err != nil {
		return fmt.Errorf("%s %q", err, arn)
	}
	if job == nil {
		d.SetId("")
		return nil
	}
	d.Set("arn", job.JobDefinitionArn)
	d.Set("container_properties", job.ContainerProperties)
	d.Set("parameters", aws.StringValueMap(job.Parameters))

	if err := d.Set("retry_strategy", flattenBatchRetryStrategy(job.RetryStrategy)); err != nil {
		return fmt.Errorf("error setting retry_strategy: %s", err)
	}

	if err := d.Set("timeout", flattenBatchJobTimeout(job.Timeout)); err != nil {
		return fmt.Errorf("error setting timeout: %s", err)
	}

	d.Set("revision", job.Revision)
	d.Set("type", job.Type)
	return nil
}

func resourceAwsBatchJobDefinitionDelete(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).batchconn
	arn := d.Get("arn").(string)
	_, err := conn.DeregisterJobDefinition(&batch.DeregisterJobDefinitionInput{
		JobDefinition: aws.String(arn),
	})
	if err != nil {
		return fmt.Errorf("%s %q", err, arn)
	}

	return nil
}

func getJobDefinition(conn *batch.Batch, arn string) (*batch.JobDefinition, error) {
	describeOpts := &batch.DescribeJobDefinitionsInput{
		JobDefinitions: []*string{aws.String(arn)},
	}
	resp, err := conn.DescribeJobDefinitions(describeOpts)
	if err != nil {
		return nil, err
	}

	numJobDefinitions := len(resp.JobDefinitions)
	switch {
	case numJobDefinitions == 0:
		return nil, nil
	case numJobDefinitions == 1:
		if *resp.JobDefinitions[0].Status == "ACTIVE" {
			return resp.JobDefinitions[0], nil
		}
		return nil, nil
	case numJobDefinitions > 1:
		return nil, fmt.Errorf("Multiple Job Definitions with name %s", arn)
	}
	return nil, nil
}

func validateAwsBatchJobContainerProperties(v interface{}, k string) (ws []string, errors []error) {
	value := v.(string)
	_, err := expandBatchJobContainerProperties(value)
	if err != nil {
		errors = append(errors, fmt.Errorf("AWS Batch Job container_properties is invalid: %s", err))
	}
	return
}

func expandBatchJobContainerProperties(rawProps string) (*batch.ContainerProperties, error) {
	var props *batch.ContainerProperties

	err := json.Unmarshal([]byte(rawProps), &props)
	if err != nil {
		return nil, fmt.Errorf("Error decoding JSON: %s", err)
	}

	return props, nil
}

func expandJobDefinitionParameters(params map[string]interface{}) map[string]*string {
	var jobParams = make(map[string]*string)
	for k, v := range params {
		jobParams[k] = aws.String(v.(string))
	}

	return jobParams
}

func expandJobDefinitionRetryStrategy(item []interface{}) *batch.RetryStrategy {
	retryStrategy := &batch.RetryStrategy{}
	data := item[0].(map[string]interface{})

	if v, ok := data["attempts"].(int); ok && v > 0 && v <= 10 {
		retryStrategy.Attempts = aws.Int64(int64(v))
	}

	return retryStrategy
}

func flattenBatchRetryStrategy(item *batch.RetryStrategy) []map[string]interface{} {
	data := []map[string]interface{}{}
	if item != nil && item.Attempts != nil {
		data = append(data, map[string]interface{}{
			"attempts": int(aws.Int64Value(item.Attempts)),
		})
	}
	return data
}

func expandJobDefinitionTimeout(item []interface{}) *batch.JobTimeout {
	timeout := &batch.JobTimeout{}
	data := item[0].(map[string]interface{})

	if v, ok := data["attempt_duration_seconds"].(int); ok && v >= 60 {
		timeout.AttemptDurationSeconds = aws.Int64(int64(v))
	}

	return timeout
}

func flattenBatchJobTimeout(item *batch.JobTimeout) []map[string]interface{} {
	data := []map[string]interface{}{}
	if item != nil && item.AttemptDurationSeconds != nil {
		data = append(data, map[string]interface{}{
			"attempt_duration_seconds": int(aws.Int64Value(item.AttemptDurationSeconds)),
		})
	}
	return data
}
<commit_msg>Fix force new behavior for aws_batch_job_definition attempts<commit_after>package aws

import (
	"fmt"

	"encoding/json"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/batch"
	"github.com/hashicorp/terraform/helper/schema"
	"github.com/hashicorp/terraform/helper/structure"
	"github.com/hashicorp/terraform/helper/validation"
)

func resourceAwsBatchJobDefinition() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsBatchJobDefinitionCreate,
		Read:   resourceAwsBatchJobDefinitionRead,
		Delete: resourceAwsBatchJobDefinitionDelete,

		Schema: map[string]*schema.Schema{
			"name": {
				Type:         schema.TypeString,
				Required:     true,
				ForceNew:     true,
				ValidateFunc: validateBatchName,
			},
			"container_properties": {
				Type:     schema.TypeString,
				Optional: true,
				ForceNew: true,
				StateFunc: func(v interface{}) string {
					json, _ := structure.NormalizeJsonString(v)
					return json
				},
				DiffSuppressFunc: suppressEquivalentJsonDiffs,
				ValidateFunc:     validateAwsBatchJobContainerProperties,
			},
			"parameters": {
				Type:     schema.TypeMap,
				Optional: true,
				ForceNew: true,
				Elem:     schema.TypeString,
			},
			"retry_strategy": {
				Type:     schema.TypeList,
				Optional: true,
				ForceNew: true,
				MaxItems: 1,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"attempts": {
							Type:         schema.TypeInt,
							Optional:     true,
							ForceNew:     true,
							ValidateFunc: validation.IntBetween(1, 10),
						},
					},
				},
			},
			"timeout": {
				Type:     schema.TypeList,
				Optional: true,
				ForceNew: true,
				MaxItems: 1,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"attempt_duration_seconds": {
							Type:         schema.TypeInt,
							Optional:     true,
							ForceNew:     true,
							ValidateFunc: validation.IntAtLeast(60),
						},
					},
				},
			},
			"type": {
				Type:         schema.TypeString,
				Required:     true,
				ForceNew:     true,
				ValidateFunc: validation.StringInSlice([]string{batch.JobDefinitionTypeContainer}, true),
			},
			"revision": {
				Type:     schema.TypeInt,
				Computed: true,
			},
			"arn": {
				Type:     schema.TypeString,
				Computed: true,
			},
		},
	}
}

func resourceAwsBatchJobDefinitionCreate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).batchconn
	name := d.Get("name").(string)

	input := &batch.RegisterJobDefinitionInput{
		JobDefinitionName: aws.String(name),
		Type:              aws.String(d.Get("type").(string)),
	}

	if v, ok := d.GetOk("container_properties"); ok {
		props, err := expandBatchJobContainerProperties(v.(string))
		if err != nil {
			return fmt.Errorf("%s %q", err, name)
		}
		input.ContainerProperties = props
	}

	if v, ok := d.GetOk("parameters"); ok {
		input.Parameters = expandJobDefinitionParameters(v.(map[string]interface{}))
	}

	if v, ok := d.GetOk("retry_strategy"); ok {
		input.RetryStrategy = expandJobDefinitionRetryStrategy(v.([]interface{}))
	}

	if v, ok := d.GetOk("timeout"); ok {
		input.Timeout = expandJobDefinitionTimeout(v.([]interface{}))
	}

	out, err := conn.RegisterJobDefinition(input)
	if err != nil {
		return fmt.Errorf("%s %q", err, name)
	}
	d.SetId(*out.JobDefinitionArn)
	d.Set("arn", out.JobDefinitionArn)
	return resourceAwsBatchJobDefinitionRead(d, meta)
}

func resourceAwsBatchJobDefinitionRead(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).batchconn
	arn := d.Get("arn").(string)
	job, err := getJobDefinition(conn, arn)
	if err != nil {
		return fmt.Errorf("%s %q", err, arn)
	}
	if job == nil {
		d.SetId("")
		return nil
	}
	d.Set("arn", job.JobDefinitionArn)
	d.Set("container_properties", job.ContainerProperties)
	d.Set("parameters", aws.StringValueMap(job.Parameters))

	if err := d.Set("retry_strategy", flattenBatchRetryStrategy(job.RetryStrategy)); err != nil {
		return fmt.Errorf("error setting retry_strategy: %s", err)
	}

	if err := d.Set("timeout", flattenBatchJobTimeout(job.Timeout)); err != nil {
		return fmt.Errorf("error setting timeout: %s", err)
	}

	d.Set("revision", job.Revision)
	d.Set("type", job.Type)
	return nil
}

func resourceAwsBatchJobDefinitionDelete(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).batchconn
	arn := d.Get("arn").(string)
	_, err := conn.DeregisterJobDefinition(&batch.DeregisterJobDefinitionInput{
		JobDefinition: aws.String(arn),
	})
	if err != nil {
		return fmt.Errorf("%s %q", err, arn)
	}

	return nil
}

func getJobDefinition(conn *batch.Batch, arn string) (*batch.JobDefinition, error) {
	describeOpts := &batch.DescribeJobDefinitionsInput{
		JobDefinitions: []*string{aws.String(arn)},
	}
	resp, err := conn.DescribeJobDefinitions(describeOpts)
	if err != nil {
		return nil, err
	}

	numJobDefinitions := len(resp.JobDefinitions)
	switch {
	case numJobDefinitions == 0:
		return nil, nil
	case numJobDefinitions == 1:
		if *resp.JobDefinitions[0].Status == "ACTIVE" {
			return resp.JobDefinitions[0], nil
		}
		return nil, nil
	case numJobDefinitions > 1:
		return nil, fmt.Errorf("Multiple Job Definitions with name %s", arn)
	}
	return nil, nil
}

func validateAwsBatchJobContainerProperties(v interface{}, k string) (ws []string, errors []error) {
	value := v.(string)
	_, err := expandBatchJobContainerProperties(value)
	if err != nil {
		errors = append(errors, fmt.Errorf("AWS Batch Job container_properties is invalid: %s", err))
	}
	return
}

func expandBatchJobContainerProperties(rawProps string) (*batch.ContainerProperties, error) {
	var props *batch.ContainerProperties

	err := json.Unmarshal([]byte(rawProps), &props)
	if err != nil {
		return nil, fmt.Errorf("Error decoding JSON: %s", err)
	}

	return props, nil
}

func expandJobDefinitionParameters(params map[string]interface{}) map[string]*string {
	var jobParams = make(map[string]*string)
	for k, v := range params {
		jobParams[k] = aws.String(v.(string))
	}

	return jobParams
}

func expandJobDefinitionRetryStrategy(item []interface{}) *batch.RetryStrategy {
	retryStrategy := &batch.RetryStrategy{}
	data := item[0].(map[string]interface{})

	if v, ok := data["attempts"].(int); ok && v > 0 && v <= 10 {
		retryStrategy.Attempts = aws.Int64(int64(v))
	}

	return retryStrategy
}

func flattenBatchRetryStrategy(item *batch.RetryStrategy) []map[string]interface{} {
	data := []map[string]interface{}{}
	if item != nil && item.Attempts != nil {
		data = append(data, map[string]interface{}{
			"attempts": int(aws.Int64Value(item.Attempts)),
		})
	}
	return data
}

func expandJobDefinitionTimeout(item []interface{}) *batch.JobTimeout {
	timeout := &batch.JobTimeout{}
	data := item[0].(map[string]interface{})

	if v, ok := data["attempt_duration_seconds"].(int); ok && v >= 60 {
		timeout.AttemptDurationSeconds = aws.Int64(int64(v))
	}

	return timeout
}

func flattenBatchJobTimeout(item *batch.JobTimeout) []map[string]interface{} {
	data := []map[string]interface{}{}
	if item != nil && item.AttemptDurationSeconds != nil {
		data = append(data, map[string]interface{}{
			"attempt_duration_seconds": int(aws.Int64Value(item.AttemptDurationSeconds)),
		})
	}
	return data
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	kafka "github.com/Shopify/sarama"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"syscall"
	"time"
	"math/rand"
)

var (
	sig_chan        = make(chan os.Signal)
	clientKill_chan = make(chan bool, 1)
	brokers         = []string{}
	topic           *string
	msgSize		*int
	clientWorkers   *int
	sentCounter     int
	chars = []rune("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ1234567890!@#$%^&*(){}][:<>.")
)

func init() {
	flag_brokers := flag.String("brokers", "localhost:9092", "Comma delimited list of Kafka brokers")
	topic = flag.String("topic", "sangrenel", "Topic to publish to")
	msgSize = flag.Int("size", 300, "Message size in bytes")
	clientWorkers = flag.Int("workers", 1, "Number of Kafka client workers")
	flag.Parse()
	brokers = strings.Split(*flag_brokers, ",")
	rand.Seed(time.Now().UnixNano())
}

func randMsg(n int) string {
    s := make([]rune, n)
    for i := range s {
        s[i] = chars[rand.Intn(len(chars))]
    }
    return string(s)
}

func sendWorker(c kafka.Client) {
	producer, err := kafka.NewProducer(&c, nil)
	if err != nil {
		panic(err)
	}
	defer producer.Close()
	for {
		err = producer.SendMessage(*topic, nil, kafka.StringEncoder(randMsg(*msgSize)))
		if err != nil {
			panic(err)
		} else {
			sentCounter++
		}
	}
}

func createClient(n int) {
	cId := "client_" + strconv.Itoa(n)
	client, err := kafka.NewClient(cId, brokers, kafka.NewClientConfig())
	if err != nil {
		panic(err)
	} else {
		fmt.Printf("%s connected\n", cId)
	}

	for i := 0; i < 5; i++ {
		go sendWorker(*client)
	}
	<-clientKill_chan
	fmt.Printf("%s shutting down\n", cId)
	client.Close()
}

func main() {
	signal.Notify(sig_chan, syscall.SIGINT, syscall.SIGTERM)
	fmt.Printf("\n::: Sangrenel :::\nStarting %s workers\nMessage size %s bytes\n\n", strconv.Itoa(*clientWorkers), strconv.Itoa(*msgSize))
	for i := 0; i < *clientWorkers; i++ {
		go createClient(i + 1)
	}
	tick := time.Tick(5 * time.Second)
	for {
		select {
		case <-tick:
			fmt.Printf("%d messages/sec published to topic: %s\n", sentCounter/5, *topic)
			sentCounter = 0
		case <-sig_chan:
			fmt.Println()
			for i := 0; i < *clientWorkers; i++ {
				clientKill_chan <- true
			}
			time.Sleep(3 * time.Second)
			os.Exit(0)
		}
	}
}
<commit_msg>useful worker error instead of grease fire<commit_after>package main

import (
	"flag"
	"fmt"
	kafka "github.com/Shopify/sarama"
	"math/rand"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"syscall"
	"time"
)

var (
	sig_chan        = make(chan os.Signal)
	clientKill_chan = make(chan bool, 1)
	brokers         = []string{}
	topic           *string
	msgSize         *int
	clientWorkers   *int
	sentCounter     int
	chars           = []rune("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ1234567890!@#$%^&*(){}][:<>.")
)

func init() {
	flag_brokers := flag.String("brokers", "localhost:9092", "Comma delimited list of Kafka brokers")
	topic = flag.String("topic", "sangrenel", "Topic to publish to")
	msgSize = flag.Int("size", 300, "Message size in bytes")
	clientWorkers = flag.Int("workers", 1, "Number of Kafka client workers")
	flag.Parse()
	brokers = strings.Split(*flag_brokers, ",")
	rand.Seed(time.Now().UnixNano())
}

func randMsg(n int) string {
	s := make([]rune, n)
	for i := range s {
		s[i] = chars[rand.Intn(len(chars))]
	}
	return string(s)
}

func sendWorker(c kafka.Client) {
	producer, err := kafka.NewProducer(&c, nil)
	if err != nil {
		panic(err)
	}
	defer producer.Close()
	for {
		err = producer.SendMessage(*topic, nil, kafka.StringEncoder(randMsg(*msgSize)))
		if err != nil {
			fmt.Println(err.Error())
		} else {
			sentCounter++
		}
	}
}

func createClient(n int) {
	cId := "client_" + strconv.Itoa(n)
	client, err := kafka.NewClient(cId, brokers, kafka.NewClientConfig())
	if err != nil {
		panic(err)
	} else {
		fmt.Printf("%s connected\n", cId)
	}

	for i := 0; i < 5; i++ {
		go sendWorker(*client)
	}
	<-clientKill_chan
	fmt.Printf("%s shutting down\n", cId)
	client.Close()
}

func main() {
	signal.Notify(sig_chan, syscall.SIGINT, syscall.SIGTERM)
	fmt.Printf("\n::: Sangrenel :::\nStarting %s workers\nMessage size %s bytes\n\n", strconv.Itoa(*clientWorkers), strconv.Itoa(*msgSize))
	for i := 0; i < *clientWorkers; i++ {
		go createClient(i + 1)
	}
	tick := time.Tick(5 * time.Second)
	for {
		select {
		case <-tick:
			fmt.Printf("%d messages/sec published to topic: %s\n", sentCounter/5, *topic)
			sentCounter = 0
		case <-sig_chan:
			fmt.Println()
			for i := 0; i < *clientWorkers; i++ {
				clientKill_chan <- true
			}
			time.Sleep(3 * time.Second)
			os.Exit(0)
		}
	}
}
<|endoftext|>


<commit_before>package sanitiser

import (
	"fmt"
	"reflect"
	"strings"
)

func parseTag(tag string) []string {

	return strings.Split(tag, ",")
}

func contains(contexts []string, context string) bool {

	for _, d := range contexts {

		if d == context {

			return true
		}
	}

	return false
}

func traverseObjects(obj interface{}, context string, hierarchy string) error {

	var v reflect.Value
	var t reflect.Type
	var ok bool

	fmt.Printf("%v.%v(type %T)\n", hierarchy, obj, obj)

	// make sure this is a pointer, so that we can update the contents if needed
	if v, ok = obj.(reflect.Value); !ok {

		v = reflect.ValueOf(obj)
	}

	for v.Kind() == reflect.Ptr || v.Kind() == reflect.Interface {

		fmt.Println("object is a pointer or an interface, calling Elem()")
		v = v.Elem()
	}

	if !v.IsValid() {

		return nil
	}

	fmt.Printf("%v.%v(type %T)\n", hierarchy, v, v)

	t = reflect.TypeOf(v.Interface())
	k := t.Kind()

	if k == reflect.Map {

		keys := v.MapKeys()
		for _, key := range keys {

			if err := traverseObjects(v.MapIndex(key), context, hierarchy+"["+fmt.Sprintf("%v", key.Elem())+"]"); err != nil {

				return err
			}
		}
	} else if k == reflect.Struct {

		for i := 0; i < t.NumField(); i++ {

			fmt.Printf("Processing field %v.%v(%v)\n", hierarchy, t.Field(i).Name, t.Field(i).Type)
			field := t.Field(i)
			field_kind := field.Type.Kind()

			if tag := field.Tag.Get("sanitise"); len(tag) > 0 {

				// the sanitise tag's value should be a comma-separated list of
				// contexts
				fmt.Printf("Field %v.%v(type %T) has a sanitise tag\n", hierarchy, field.Name, v.Field(i))
				contexts := parseTag(tag)
				if contains(contexts, context) || contains(contexts, "*") {
					// sanitise this field
					if !v.Field(i).CanSet() {

						return fmt.Errorf("Unable to set zero value for %v.%v", hierarchy, t.Field(i).Name)
					}

					fmt.Printf("Sanitising field %v.%v\n", hierarchy, t.Field(i).Name)
					v.Field(i).Set(reflect.New(t.Field(i).Type).Elem())
				}
			} else if field_kind == reflect.Struct || field_kind == reflect.Interface {

				sv := v.Field(i)
				fmt.Printf("Processing object %v.%v(type %T)\n", hierarchy, sv, sv)

				if err := traverseObjects(sv, context, hierarchy+"."+t.Field(i).Name); err != nil {

					return err
				}
			}
		}
	}

	return nil
}

func Sanitise(obj interface{}, context string) error {

	return traverseObjects(obj, context, "")
}
<commit_msg>Added an option to set a logging function<commit_after>package sanitiser

import (
	"fmt"
	"reflect"
	"strings"
)

type Logger func(format string, v ...interface{})

var dbg Logger

func init() {

	dbg = func(string, ...interface{}) {}
}

func parseTag(tag string) []string {

	return strings.Split(tag, ",")
}

func contains(contexts []string, context string) bool {

	for _, d := range contexts {

		if d == context {

			return true
		}
	}

	return false
}

func traverseObjects(obj interface{}, context string, hierarchy string) error {

	var v reflect.Value
	var t reflect.Type
	var ok bool

	dbg("%v.%v(type %T)\n", hierarchy, obj, obj)

	// make sure this is a pointer, so that we can update the contents if needed
	if v, ok = obj.(reflect.Value); !ok {

		v = reflect.ValueOf(obj)
	}

	for v.Kind() == reflect.Ptr || v.Kind() == reflect.Interface {

		dbg("object is a pointer or an interface, calling Elem()\n")
		v = v.Elem()
	}

	if !v.IsValid() {

		return nil
	}

	dbg("%v.%v(type %T)\n", hierarchy, v, v)

	t = reflect.TypeOf(v.Interface())
	k := t.Kind()

	if k == reflect.Map {

		keys := v.MapKeys()
		for _, key := range keys {

			if err := traverseObjects(v.MapIndex(key), context, hierarchy+"["+fmt.Sprintf("%v", key.Elem())+"]"); err != nil {

				return err
			}
		}
	} else if k == reflect.Struct {

		for i := 0; i < t.NumField(); i++ {

			dbg("Processing field %v.%v(%v)\n", hierarchy, t.Field(i).Name, t.Field(i).Type)
			field := t.Field(i)
			field_kind := field.Type.Kind()

			if tag := field.Tag.Get("sanitise"); len(tag) > 0 {

				// the sanitise tag's value should be a comma-separated list of
				// contexts
				dbg("Field %v.%v(type %T) has a sanitise tag\n", hierarchy, field.Name, v.Field(i))
				contexts := parseTag(tag)
				if contains(contexts, context) || contains(contexts, "*") {
					// sanitise this field
					if !v.Field(i).CanSet() {

						return fmt.Errorf("Unable to set zero value for %v.%v", hierarchy, t.Field(i).Name)
					}

					dbg("Sanitising field %v.%v\n", hierarchy, t.Field(i).Name)
					v.Field(i).Set(reflect.New(t.Field(i).Type).Elem())
				}
			} else if field_kind == reflect.Struct || field_kind == reflect.Interface {

				sv := v.Field(i)
				dbg("Processing object %v.%v(type %T)\n", hierarchy, sv, sv)

				if err := traverseObjects(sv, context, hierarchy+"."+t.Field(i).Name); err != nil {

					return err
				}
			}
		}
	}

	return nil
}

func Sanitise(obj interface{}, context string) error {

	return traverseObjects(obj, context, "")
}

func SetLogger(f Logger) {

	dbg = f
}
<|endoftext|>


<commit_before>package storm

import (
	"sync"
	"testing"
	"time"

	"github.com/asdine/storm/codec/gob"
	"github.com/asdine/storm/codec/json"
	"github.com/boltdb/bolt"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestSave(t *testing.T) {
	db, cleanup := createDB(t)
	defer cleanup()

	err := db.Save(&SimpleUser{ID: 10, Name: "John"})
	assert.NoError(t, err)

	err = db.Save(&SimpleUser{Name: "John"})
	assert.Error(t, err)
	assert.Equal(t, ErrZeroID, err)

	err = db.Save(&ClassicBadTags{ID: "id", PublicField: 100})
	assert.Error(t, err)
	assert.Equal(t, ErrUnknownTag, err)

	err = db.Save(&UserWithNoID{Name: "John"})
	assert.Error(t, err)
	assert.Equal(t, ErrNoID, err)

	err = db.Save(&UserWithIDField{ID: 10, Name: "John"})
	assert.NoError(t, err)

	u := UserWithEmbeddedIDField{}
	u.ID = 150
	u.Name = "Pete"
	u.Age = 10
	err = db.Save(&u)
	assert.NoError(t, err)

	v := UserWithIDField{ID: 10, Name: "John"}
	err = db.Save(&v)
	assert.NoError(t, err)

	w := UserWithEmbeddedField{}
	w.ID = 150
	w.Name = "John"
	err = db.Save(&w)
	assert.NoError(t, err)

	db.Bolt.View(func(tx *bolt.Tx) error {
		bucket := tx.Bucket([]byte("UserWithIDField"))
		assert.NotNil(t, bucket)

		i, err := toBytes(10, json.Codec)
		assert.NoError(t, err)

		val := bucket.Get(i)
		assert.NotNil(t, val)

		content, err := db.codec.Marshal(&v)
		assert.NoError(t, err)
		assert.Equal(t, content, val)
		return nil
	})
}

func TestSaveUnique(t *testing.T) {
	db, cleanup := createDB(t)
	defer cleanup()

	u1 := UniqueNameUser{ID: 10, Name: "John", age: 10}
	err := db.Save(&u1)
	assert.NoError(t, err)

	u2 := UniqueNameUser{ID: 11, Name: "John", age: 100}
	err = db.Save(&u2)
	assert.Error(t, err)
	assert.True(t, ErrAlreadyExists == err)

	// same id
	u3 := UniqueNameUser{ID: 10, Name: "Jake", age: 100}
	err = db.Save(&u3)
	assert.NoError(t, err)

	db.Bolt.View(func(tx *bolt.Tx) error {
		bucket := tx.Bucket([]byte("UniqueNameUser"))

		uniqueBucket := bucket.Bucket([]byte(indexPrefix + "Name"))
		assert.NotNil(t, uniqueBucket)

		id := uniqueBucket.Get([]byte("Jake"))
		i, err := toBytes(10, json.Codec)
		assert.NoError(t, err)
		assert.Equal(t, i, id)

		id = uniqueBucket.Get([]byte("John"))
		assert.Nil(t, id)
		return nil
	})
}

func TestSaveUniqueStruct(t *testing.T) {
	db, cleanup := createDB(t)
	defer cleanup()

	a := ClassicUnique{ID: "id1"}
	a.InlineStruct.A = 10.0
	a.InlineStruct.B = 12.0

	err := db.Save(&a)
	require.NoError(t, err)

	b := ClassicUnique{ID: "id2"}
	b.InlineStruct.A = 10.0
	b.InlineStruct.B = 12.0

	err = db.Save(&b)
	require.Equal(t, ErrAlreadyExists, err)
}

func TestSaveIndex(t *testing.T) {
	db, cleanup := createDB(t)
	defer cleanup()

	u1 := IndexedNameUser{ID: 10, Name: "John", age: 10}
	err := db.Save(&u1)
	assert.NoError(t, err)

	u1 = IndexedNameUser{ID: 10, Name: "John", age: 10}
	err = db.Save(&u1)
	assert.NoError(t, err)

	u2 := IndexedNameUser{ID: 11, Name: "John", age: 100}
	err = db.Save(&u2)
	assert.NoError(t, err)

	name1 := "Jake"
	name2 := "Jane"
	name3 := "James"

	for i := 0; i < 100; i++ {
		u := IndexedNameUser{ID: i + 1}

		if i%2 == 0 {
			u.Name = name1
		} else {
			u.Name = name2
		}

		db.Save(&u)
	}

	var users []IndexedNameUser
	err = db.Find("Name", name1, &users)
	assert.NoError(t, err)
	assert.Len(t, users, 50)

	err = db.Find("Name", name2, &users)
	assert.NoError(t, err)
	assert.Len(t, users, 50)

	err = db.Find("Name", name3, &users)
	assert.Error(t, err)
	assert.Equal(t, ErrNotFound, err)

	err = db.Save(nil)
	assert.Error(t, err)
	assert.Equal(t, ErrStructPtrNeeded, err)
}

func TestSaveEmptyValues(t *testing.T) {
	db, cleanup := createDB(t)
	defer cleanup()

	u := User{
		ID: 10,
	}
	err := db.Save(&u)
	assert.NoError(t, err)

	var v User
	err = db.One("ID", 10, &v)
	assert.NoError(t, err)
	assert.Equal(t, 10, v.ID)

	u.Name = "John"
	u.Slug = "john"
	err = db.Save(&u)
	assert.NoError(t, err)

	err = db.One("Name", "John", &v)
	assert.NoError(t, err)
	assert.Equal(t, "John", v.Name)
	assert.Equal(t, "john", v.Slug)
	err = db.One("Slug", "john", &v)
	assert.NoError(t, err)
	assert.Equal(t, "John", v.Name)
	assert.Equal(t, "john", v.Slug)

	u.Name = ""
	u.Slug = ""
	err = db.Save(&u)
	assert.NoError(t, err)

	err = db.One("Name", "John", &v)
	assert.Error(t, err)
	err = db.One("Slug", "john", &v)
	assert.Error(t, err)
}

func TestSaveAutoIncrement(t *testing.T) {
	db, cleanup := createDB(t, AutoIncrement())
	defer cleanup()

	for i := 1; i < 10; i++ {
		s := SimpleUser{Name: "John"}
		err := db.Save(&s)
		assert.NoError(t, err)
		assert.Equal(t, i, s.ID)
	}

	u := UserWithUint64IDField{Name: "John"}
	err := db.Save(&u)
	assert.NoError(t, err)
	assert.Equal(t, uint64(1), u.ID)
	v := UserWithUint64IDField{}
	err = db.One("ID", uint64(1), &v)
	assert.NoError(t, err)
	assert.Equal(t, u, v)

	ui := UserWithIDField{Name: "John"}
	err = db.Save(&ui)
	assert.NoError(t, err)
	assert.Equal(t, 1, ui.ID)
	vi := UserWithIDField{}
	err = db.One("ID", 1, &vi)
	assert.NoError(t, err)
	assert.Equal(t, ui, vi)

	us := UserWithStringIDField{Name: "John"}
	err = db.Save(&us)
	assert.Error(t, err)
	assert.Equal(t, ErrZeroID, err)
}

func TestSaveDifferentBucketRoot(t *testing.T) {
	db, cleanup := createDB(t)
	defer cleanup()

	assert.Len(t, db.rootBucket, 0)

	dbSub := db.From("sub").(*node)

	assert.NotEqual(t, dbSub, db)
	assert.Len(t, dbSub.rootBucket, 1)

	err := db.Save(&User{ID: 10, Name: "John"})
	assert.NoError(t, err)
	err = dbSub.Save(&User{ID: 11, Name: "Paul"})
	assert.NoError(t, err)

	var (
		john User
		paul User
	)

	err = db.One("Name", "John", &john)
	assert.NoError(t, err)
	err = db.One("Name", "Paul", &paul)
	assert.Error(t, err)

	err = dbSub.One("Name", "Paul", &paul)
	assert.NoError(t, err)
	err = dbSub.One("Name", "John", &john)
	assert.Error(t, err)
}

func TestSaveEmbedded(t *testing.T) {
	db, cleanup := createDB(t)
	defer cleanup()
	AutoIncrement()(db)

	type Base struct {
		ID int `storm:"id"`
	}

	type User struct {
		Base      `storm:"inline"`
		Group     string `storm:"index"`
		Email     string `storm:"unique"`
		Name      string
		Age       int
		CreatedAt time.Time `storm:"index"`
	}

	user := User{
		Group:     "staff",
		Email:     "john@provider.com",
		Name:      "John",
		Age:       21,
		CreatedAt: time.Now(),
	}

	err := db.Save(&user)
	assert.NoError(t, err)
	assert.Equal(t, 1, user.ID)
}

func TestSaveByValue(t *testing.T) {
	db, cleanup := createDB(t)
	defer cleanup()

	w := User{Name: "John"}
	err := db.Save(w)
	assert.Error(t, err)
	assert.Equal(t, ErrStructPtrNeeded, err)
}

func TestSaveWithBatch(t *testing.T) {
	db, cleanup := createDB(t, Batch())
	defer cleanup()

	var wg sync.WaitGroup

	for i := 0; i < 5; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			err := db.Save(&User{ID: i + 1, Name: "John"})
			assert.NoError(t, err)
		}()
	}

	wg.Wait()
}

func TestSaveMetadata(t *testing.T) {
	db, cleanup := createDB(t, Batch())
	defer cleanup()

	w := User{ID: 10, Name: "John"}
	err := db.Save(&w)
	require.NoError(t, err)
	n := db.WithCodec(gob.Codec)
	err = n.Save(&w)
	require.Equal(t, ErrDifferentCodec, err)
}

func BenchmarkSave(b *testing.B) {
	db, cleanup := createDB(b, AutoIncrement())
	defer cleanup()

	w := User{Name: "John"}
	b.ResetTimer()
	for n := 0; n < b.N; n++ {
		err := db.Save(&w)
		if err != nil {
			b.Error(err)
		}
	}
}
<commit_msg>Test struct index with One<commit_after>package storm

import (
	"sync"
	"testing"
	"time"

	"github.com/asdine/storm/codec/gob"
	"github.com/asdine/storm/codec/json"
	"github.com/boltdb/bolt"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestSave(t *testing.T) {
	db, cleanup := createDB(t)
	defer cleanup()

	err := db.Save(&SimpleUser{ID: 10, Name: "John"})
	assert.NoError(t, err)

	err = db.Save(&SimpleUser{Name: "John"})
	assert.Error(t, err)
	assert.Equal(t, ErrZeroID, err)

	err = db.Save(&ClassicBadTags{ID: "id", PublicField: 100})
	assert.Error(t, err)
	assert.Equal(t, ErrUnknownTag, err)

	err = db.Save(&UserWithNoID{Name: "John"})
	assert.Error(t, err)
	assert.Equal(t, ErrNoID, err)

	err = db.Save(&UserWithIDField{ID: 10, Name: "John"})
	assert.NoError(t, err)

	u := UserWithEmbeddedIDField{}
	u.ID = 150
	u.Name = "Pete"
	u.Age = 10
	err = db.Save(&u)
	assert.NoError(t, err)

	v := UserWithIDField{ID: 10, Name: "John"}
	err = db.Save(&v)
	assert.NoError(t, err)

	w := UserWithEmbeddedField{}
	w.ID = 150
	w.Name = "John"
	err = db.Save(&w)
	assert.NoError(t, err)

	db.Bolt.View(func(tx *bolt.Tx) error {
		bucket := tx.Bucket([]byte("UserWithIDField"))
		assert.NotNil(t, bucket)

		i, err := toBytes(10, json.Codec)
		assert.NoError(t, err)

		val := bucket.Get(i)
		assert.NotNil(t, val)

		content, err := db.codec.Marshal(&v)
		assert.NoError(t, err)
		assert.Equal(t, content, val)
		return nil
	})
}

func TestSaveUnique(t *testing.T) {
	db, cleanup := createDB(t)
	defer cleanup()

	u1 := UniqueNameUser{ID: 10, Name: "John", age: 10}
	err := db.Save(&u1)
	assert.NoError(t, err)

	u2 := UniqueNameUser{ID: 11, Name: "John", age: 100}
	err = db.Save(&u2)
	assert.Error(t, err)
	assert.True(t, ErrAlreadyExists == err)

	// same id
	u3 := UniqueNameUser{ID: 10, Name: "Jake", age: 100}
	err = db.Save(&u3)
	assert.NoError(t, err)

	db.Bolt.View(func(tx *bolt.Tx) error {
		bucket := tx.Bucket([]byte("UniqueNameUser"))

		uniqueBucket := bucket.Bucket([]byte(indexPrefix + "Name"))
		assert.NotNil(t, uniqueBucket)

		id := uniqueBucket.Get([]byte("Jake"))
		i, err := toBytes(10, json.Codec)
		assert.NoError(t, err)
		assert.Equal(t, i, id)

		id = uniqueBucket.Get([]byte("John"))
		assert.Nil(t, id)
		return nil
	})
}

func TestSaveUniqueStruct(t *testing.T) {
	db, cleanup := createDB(t)
	defer cleanup()

	a := ClassicUnique{ID: "id1"}
	a.InlineStruct.A = 10.0
	a.InlineStruct.B = 12.0

	err := db.Save(&a)
	require.NoError(t, err)

	b := ClassicUnique{ID: "id2"}
	b.InlineStruct.A = 10.0
	b.InlineStruct.B = 12.0

	err = db.Save(&b)
	require.Equal(t, ErrAlreadyExists, err)

	err = db.One("InlineStruct", struct {
		A float32
		B float64
	}{A: 10.0, B: 12.0}, &b)
	require.NoError(t, err)
	require.Equal(t, a.ID, b.ID)
}

func TestSaveIndex(t *testing.T) {
	db, cleanup := createDB(t)
	defer cleanup()

	u1 := IndexedNameUser{ID: 10, Name: "John", age: 10}
	err := db.Save(&u1)
	assert.NoError(t, err)

	u1 = IndexedNameUser{ID: 10, Name: "John", age: 10}
	err = db.Save(&u1)
	assert.NoError(t, err)

	u2 := IndexedNameUser{ID: 11, Name: "John", age: 100}
	err = db.Save(&u2)
	assert.NoError(t, err)

	name1 := "Jake"
	name2 := "Jane"
	name3 := "James"

	for i := 0; i < 100; i++ {
		u := IndexedNameUser{ID: i + 1}

		if i%2 == 0 {
			u.Name = name1
		} else {
			u.Name = name2
		}

		db.Save(&u)
	}

	var users []IndexedNameUser
	err = db.Find("Name", name1, &users)
	assert.NoError(t, err)
	assert.Len(t, users, 50)

	err = db.Find("Name", name2, &users)
	assert.NoError(t, err)
	assert.Len(t, users, 50)

	err = db.Find("Name", name3, &users)
	assert.Error(t, err)
	assert.Equal(t, ErrNotFound, err)

	err = db.Save(nil)
	assert.Error(t, err)
	assert.Equal(t, ErrStructPtrNeeded, err)
}

func TestSaveEmptyValues(t *testing.T) {
	db, cleanup := createDB(t)
	defer cleanup()

	u := User{
		ID: 10,
	}
	err := db.Save(&u)
	assert.NoError(t, err)

	var v User
	err = db.One("ID", 10, &v)
	assert.NoError(t, err)
	assert.Equal(t, 10, v.ID)

	u.Name = "John"
	u.Slug = "john"
	err = db.Save(&u)
	assert.NoError(t, err)

	err = db.One("Name", "John", &v)
	assert.NoError(t, err)
	assert.Equal(t, "John", v.Name)
	assert.Equal(t, "john", v.Slug)
	err = db.One("Slug", "john", &v)
	assert.NoError(t, err)
	assert.Equal(t, "John", v.Name)
	assert.Equal(t, "john", v.Slug)

	u.Name = ""
	u.Slug = ""
	err = db.Save(&u)
	assert.NoError(t, err)

	err = db.One("Name", "John", &v)
	assert.Error(t, err)
	err = db.One("Slug", "john", &v)
	assert.Error(t, err)
}

func TestSaveAutoIncrement(t *testing.T) {
	db, cleanup := createDB(t, AutoIncrement())
	defer cleanup()

	for i := 1; i < 10; i++ {
		s := SimpleUser{Name: "John"}
		err := db.Save(&s)
		assert.NoError(t, err)
		assert.Equal(t, i, s.ID)
	}

	u := UserWithUint64IDField{Name: "John"}
	err := db.Save(&u)
	assert.NoError(t, err)
	assert.Equal(t, uint64(1), u.ID)
	v := UserWithUint64IDField{}
	err = db.One("ID", uint64(1), &v)
	assert.NoError(t, err)
	assert.Equal(t, u, v)

	ui := UserWithIDField{Name: "John"}
	err = db.Save(&ui)
	assert.NoError(t, err)
	assert.Equal(t, 1, ui.ID)
	vi := UserWithIDField{}
	err = db.One("ID", 1, &vi)
	assert.NoError(t, err)
	assert.Equal(t, ui, vi)

	us := UserWithStringIDField{Name: "John"}
	err = db.Save(&us)
	assert.Error(t, err)
	assert.Equal(t, ErrZeroID, err)
}

func TestSaveDifferentBucketRoot(t *testing.T) {
	db, cleanup := createDB(t)
	defer cleanup()

	assert.Len(t, db.rootBucket, 0)

	dbSub := db.From("sub").(*node)

	assert.NotEqual(t, dbSub, db)
	assert.Len(t, dbSub.rootBucket, 1)

	err := db.Save(&User{ID: 10, Name: "John"})
	assert.NoError(t, err)
	err = dbSub.Save(&User{ID: 11, Name: "Paul"})
	assert.NoError(t, err)

	var (
		john User
		paul User
	)

	err = db.One("Name", "John", &john)
	assert.NoError(t, err)
	err = db.One("Name", "Paul", &paul)
	assert.Error(t, err)

	err = dbSub.One("Name", "Paul", &paul)
	assert.NoError(t, err)
	err = dbSub.One("Name", "John", &john)
	assert.Error(t, err)
}

func TestSaveEmbedded(t *testing.T) {
	db, cleanup := createDB(t)
	defer cleanup()
	AutoIncrement()(db)

	type Base struct {
		ID int `storm:"id"`
	}

	type User struct {
		Base      `storm:"inline"`
		Group     string `storm:"index"`
		Email     string `storm:"unique"`
		Name      string
		Age       int
		CreatedAt time.Time `storm:"index"`
	}

	user := User{
		Group:     "staff",
		Email:     "john@provider.com",
		Name:      "John",
		Age:       21,
		CreatedAt: time.Now(),
	}

	err := db.Save(&user)
	assert.NoError(t, err)
	assert.Equal(t, 1, user.ID)
}

func TestSaveByValue(t *testing.T) {
	db, cleanup := createDB(t)
	defer cleanup()

	w := User{Name: "John"}
	err := db.Save(w)
	assert.Error(t, err)
	assert.Equal(t, ErrStructPtrNeeded, err)
}

func TestSaveWithBatch(t *testing.T) {
	db, cleanup := createDB(t, Batch())
	defer cleanup()

	var wg sync.WaitGroup

	for i := 0; i < 5; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			err := db.Save(&User{ID: i + 1, Name: "John"})
			assert.NoError(t, err)
		}()
	}

	wg.Wait()
}

func TestSaveMetadata(t *testing.T) {
	db, cleanup := createDB(t, Batch())
	defer cleanup()

	w := User{ID: 10, Name: "John"}
	err := db.Save(&w)
	require.NoError(t, err)
	n := db.WithCodec(gob.Codec)
	err = n.Save(&w)
	require.Equal(t, ErrDifferentCodec, err)
}

func BenchmarkSave(b *testing.B) {
	db, cleanup := createDB(b, AutoIncrement())
	defer cleanup()

	w := User{Name: "John"}
	b.ResetTimer()
	for n := 0; n < b.N; n++ {
		err := db.Save(&w)
		if err != nil {
			b.Error(err)
		}
	}
}
<|endoftext|>


<commit_before>package compiler

import (
	"github.com/davyxu/tabtoy/v3/helper"
	"github.com/davyxu/tabtoy/v3/model"
	"github.com/davyxu/tabtoy/v3/report"
	"strings"
)

func LoadHeader(sheet helper.TableSheet, tab *model.DataTable, resolveTableType string, typeTab *model.TypeTable) (maxCol int) {
	// 读取表头

	for col := 0; ; col++ {

		headerValue := sheet.GetValue(0, col, nil)

		// 空列，终止
		if headerValue == "" {
			break
		}

		maxCol = col
		// 列头带#时，本列忽略
		if strings.HasPrefix(headerValue, "#") {
			continue
		}

		header := tab.MustGetHeader(col)
		header.Cell.CopyFrom(&model.Cell{
			Value: headerValue,
			Col:   col,
			Row:   0,
			Table: tab,
		})

	}

	resolveHeaderFields(tab, resolveTableType, typeTab)

	checkHeaderTypes(tab, typeTab)

	return
}

func checkHeaderTypes(tab *model.DataTable, typeTab *model.TypeTable) {

	for _, header := range tab.Headers {

		if header.TypeInfo == nil {
			continue
		}

		// 原始类型检查
		if !model.PrimitiveExists(header.TypeInfo.FieldType) &&
			!typeTab.ObjectExists(header.TypeInfo.FieldType) { // 对象检查

			report.ReportError("UnknownFieldType", header.TypeInfo.FieldType, header.Cell.String())
		}
	}

}

func headerValueExists(offset int, name string, headers []*model.HeaderField) bool {

	for i := offset; i < len(headers); i++ {
		if headers[i].Cell.Value == name {
			return true
		}
	}

	return false
}

func resolveHeaderFields(tab *model.DataTable, tableObjectType string, typeTab *model.TypeTable) {

	tab.OriginalHeaderType = tableObjectType
	for index, header := range tab.Headers {

		if header.Cell.Value == "" {
			continue
		}

		tf := typeTab.FieldByName(tableObjectType, header.Cell.Value)
		if tf == nil {
			report.ReportError("HeaderFieldNotDefined", header.Cell.String())
		}

		if headerValueExists(index+1, header.Cell.Value, tab.Headers) && !tf.IsArray() {
			report.ReportError("DuplicateHeaderField", header.Cell.String())
		}

		// 解析好的类型
		header.TypeInfo = tf
	}

}
<commit_msg>加强头未定义时的报错<commit_after>package compiler

import (
	"github.com/davyxu/tabtoy/v3/helper"
	"github.com/davyxu/tabtoy/v3/model"
	"github.com/davyxu/tabtoy/v3/report"
	"strings"
)

func LoadHeader(sheet helper.TableSheet, tab *model.DataTable, resolveTableType string, typeTab *model.TypeTable) (maxCol int) {
	// 读取表头

	for col := 0; ; col++ {

		headerValue := sheet.GetValue(0, col, nil)

		// 空列，终止
		if headerValue == "" {
			break
		}

		maxCol = col
		// 列头带#时，本列忽略
		if strings.HasPrefix(headerValue, "#") {
			continue
		}

		header := tab.MustGetHeader(col)
		header.Cell.CopyFrom(&model.Cell{
			Value: headerValue,
			Col:   col,
			Row:   0,
			Table: tab,
		})

	}

	resolveHeaderFields(tab, resolveTableType, typeTab)

	checkHeaderTypes(tab, typeTab)

	return
}

func checkHeaderTypes(tab *model.DataTable, typeTab *model.TypeTable) {

	for _, header := range tab.Headers {

		if header.TypeInfo == nil {
			continue
		}

		// 原始类型检查
		if !model.PrimitiveExists(header.TypeInfo.FieldType) &&
			!typeTab.ObjectExists(header.TypeInfo.FieldType) { // 对象检查

			report.ReportError("UnknownFieldType", header.TypeInfo.FieldType, header.Cell.String())
		}
	}

}

func headerValueExists(offset int, name string, headers []*model.HeaderField) bool {

	for i := offset; i < len(headers); i++ {
		if headers[i].Cell.Value == name {
			return true
		}
	}

	return false
}

func resolveHeaderFields(tab *model.DataTable, tableObjectType string, typeTab *model.TypeTable) {

	tab.OriginalHeaderType = tableObjectType
	for index, header := range tab.Headers {

		if header.Cell.Value == "" {
			continue
		}

		tf := typeTab.FieldByName(tableObjectType, header.Cell.Value)
		if tf == nil {
			report.ReportError("HeaderFieldNotDefined", header.Cell.String(), tableObjectType)
		}

		if headerValueExists(index+1, header.Cell.Value, tab.Headers) && !tf.IsArray() {
			report.ReportError("DuplicateHeaderField", header.Cell.String())
		}

		// 解析好的类型
		header.TypeInfo = tf
	}

}
<|endoftext|>