<commit_before>package lnwallet

import (
	"errors"
	"fmt"
	"sync"

	"github.com/roasbeef/btcd/btcec"
	"github.com/roasbeef/btcd/chaincfg/chainhash"
	"github.com/roasbeef/btcd/txscript"
	"github.com/roasbeef/btcd/wire"
	"github.com/roasbeef/btcutil"
)

// ErrNotMine is an error denoting that a WalletController instance is unable
// to spend a specifid output.
var ErrNotMine = errors.New("the passed output doesn't belong to the wallet")

// AddressType is a enum-like type which denotes the possible address types
// WalletController supports.
type AddressType uint8

const (
	// WitnessPubKey represents a p2wkh address.
	WitnessPubKey AddressType = iota

	// NestedWitnessPubKey represents a p2sh output which is itself a
	// nested p2wkh output.
	NestedWitnessPubKey

	// PubKeyHash represents a regular p2pkh output.
	PubKeyHash
)

// Utxo is an unspent output denoted by its outpoint, and output value of the
// original output.
type Utxo struct {
	Value btcutil.Amount
	wire.OutPoint
}

// TransactionDetail describes a transaction with either inputs which belong to
// the wallet, or has outputs that pay to the wallet.
type TransactionDetail struct {
	// Hash is the transaction hash of the transaction.
	Hash chainhash.Hash

	// Value is the net value of this transaction (in satoshis) from the
	// PoV of the wallet. If this transaction purely spends from the
	// wallet's funds, then this value will be negative. Similarly, if this
	// transaction credits the wallet, then this value will be positive.
	Value btcutil.Amount

	// NumConfirmations is the number of confirmations this transaction
	// has. If the transaction is unconfirmed, then this value will be
	// zero.
	NumConfirmations int32

	// BlockHeight is the hash of the block which includes this
	// transaction. Unconfirmed transactions will have a nil value for this
	// field.
	BlockHash *chainhash.Hash

	// BlockHeight is the height of the block including this transaction.
	// Unconfirmed transaction will show a height of zero.
	BlockHeight int32

	// Timestamp is the unix timestamp of the block including this
	// transaction. If the transaction is unconfirmed, then this will be a
	// timestamp of txn creation.
	Timestamp int64

	// TotalFees is the total fee in satoshis paid by this transaction.
	TotalFees int64
}

// TransactionSubscription is an interface which describes an object capable of
// receiving notifications of new transaction related to the underlying wallet.
// TODO(roasbeef): add balance updates?
type TransactionSubscription interface {
	// ConfirmedTransactions returns a channel which will be sent on as new
	// relevant transactions are confirmed.
	ConfirmedTransactions() chan *TransactionDetail

	// UnconfirmedTransactions returns a channel which will be sent on as
	// new relevant transactions are seen within the network.
	UnconfirmedTransactions() chan *TransactionDetail

	// Cancel finalizes the subscription, cleaning up any resources
	// allocated.
	Cancel()
}

// WalletController defines an abstract interface for controlling a local Pure
// Go wallet, a local or remote wallet via an RPC mechanism, or possibly even
// a daemon assisted hardware wallet. This interface serves the purpose of
// allowing LightningWallet to be seamlessly compatible with several wallets
// such as: uspv, btcwallet, Bitcoin Core, Electrum, etc. This interface then
// serves as a "base wallet", with Lightning Network awareness taking place at
// a "higher" level of abstraction. Essentially, an overlay wallet.
// Implementors of this interface must closely adhere to the documented
// behavior of all interface methods in order to ensure identical behavior
// across all concrete implementations.
type WalletController interface {
	// FetchInputInfo queries for the WalletController's knowledge of the
	// passed outpoint. If the base wallet determines this output is under
	// its control, then the original txout should be returned. Otherwise,
	// a non-nil error value of ErrNotMine should be returned instead.
	FetchInputInfo(prevOut *wire.OutPoint) (*wire.TxOut, error)

	// ConfirmedBalance returns the sum of all the wallet's unspent outputs
	// that have at least confs confirmations. If confs is set to zero,
	// then all unspent outputs, including those currently in the mempool
	// will be included in the final sum.
	ConfirmedBalance(confs int32, witness bool) (btcutil.Amount, error)

	// NewAddress returns the next external or internal address for the
	// wallet dictated by the value of the `change` parameter. If change is
	// true, then an internal address should be used, otherwise an external
	// address should be returned. The type of address returned is dictated
	// by the wallet's capabilities, and may be of type: p2sh, p2pkh,
	// p2wkh, p2wsh, etc.
	NewAddress(addrType AddressType, change bool) (btcutil.Address, error)

	// GetPrivKey retrives the underlying private key associated with the
	// passed address. If the wallet is unable to locate this private key
	// due to the address not being under control of the wallet, then an
	// error should be returned.
	GetPrivKey(a btcutil.Address) (*btcec.PrivateKey, error)

	// NewRawKey returns a raw private key controlled by the wallet. These
	// keys are used for the 2-of-2 multi-sig outputs for funding
	// transactions, as well as the pub key used for commitment transactions.
	//
	// NOTE: The wallet MUST watch for on-chain outputs created to a p2wpkh
	// script using keys returned by this function.
	NewRawKey() (*btcec.PublicKey, error)

	// FetchRootKey returns a root key which will be used by the
	// LightningWallet to deterministically generate secrets. The private
	// key returned by this method should remain constant in-between
	// WalletController restarts.
	FetchRootKey() (*btcec.PrivateKey, error)

	// SendOutputs funds, signs, and broadcasts a Bitcoin transaction
	// paying out to the specified outputs. In the case the wallet has
	// insufficient funds, or the outputs are non-standard, an error
	// should be returned.
	SendOutputs(outputs []*wire.TxOut) (*chainhash.Hash, error)

	// ListUnspentWitness returns all unspent outputs which are version 0
	// witness programs. The 'confirms' parameter indicates the minimum
	// number of confirmations an output needs in order to be returned by
	// this method. Passing -1 as 'confirms' indicates that even
	// unconfirmed outputs should be returned.
	ListUnspentWitness(confirms int32) ([]*Utxo, error)

	// ListTransactionDetails returns a list of all transactions which are
	// relevant to the wallet.
	ListTransactionDetails() ([]*TransactionDetail, error)

	// LockOutpoint marks an outpoint as locked meaning it will no longer
	// be deemed as eligible for coin selection. Locking outputs are
	// utilized in order to avoid race conditions when selecting inputs for
	// usage when funding a channel.
	LockOutpoint(o wire.OutPoint)

	// UnlockOutpoint unlocks an previously locked output, marking it
	// eligible for coin selection.
	UnlockOutpoint(o wire.OutPoint)

	// PublishTransaction performs cursory validation (dust checks, etc),
	// then finally broadcasts the passed transaction to the Bitcoin network.
	PublishTransaction(tx *wire.MsgTx) error

	// SubscribeTransactions returns a TransactionSubscription client which
	// is capable of receiving async notifications as new transactions
	// related to the wallet are seen within the network, or found in
	// blocks.
	//
	// NOTE: a non-nil error should be returned if notifications aren't
	// supported.
	//
	// TODO(roasbeef): make distinct interface?
	SubscribeTransactions() (TransactionSubscription, error)

	// IsSynced returns a boolean indicating if from the PoV of the wallet,
	// it has fully synced to the current best block in the main chain.
	IsSynced() (bool, error)

	// Start initializes the wallet, making any necessary connections,
	// starting up required goroutines etc.
	Start() error

	// Stop signals the wallet for shutdown. Shutdown may entail closing
	// any active sockets, database handles, stopping goroutines, etc.
	Stop() error
}

// BlockChainIO is a dedicated source which will be used to obtain queries
// related to the current state of the blockchain. The data returned by each of
// the defined methods within this interface should always return the most up
// to date data possible.
//
// TODO(roasbeef): move to diff package perhaps?
// TODO(roasbeef): move publish txn here?
type BlockChainIO interface {
	// GetBestBlock returns the current height and block hash of the valid
	// most-work chain the implementation is aware of.
	GetBestBlock() (*chainhash.Hash, int32, error)

	// GetTxOut returns the original output referenced by the passed
	// outpoint.
	GetUtxo(txid *chainhash.Hash, index uint32) (*wire.TxOut, error)

	// GetTransaction returns the full transaction identified by the passed
	// transaction ID.
	GetTransaction(txid *chainhash.Hash) (*wire.MsgTx, error)

	// GetBlockHash returns the hash of the block in the best blockchain
	// at the given height.
	GetBlockHash(blockHeight int64) (*chainhash.Hash, error)

	// GetBlock returns the block in the main chain identified by the given
	// hash.
	GetBlock(blockHash *chainhash.Hash) (*wire.MsgBlock, error)
}

// SignDescriptor houses the necessary information required to successfully sign
// a given output. This struct is used by the Signer interface in order to gain
// access to critical data needed to generate a valid signature.
type SignDescriptor struct {
	// Pubkey is the public key to which the signature should be generated
	// over. The Signer should then generate a signature with the private
	// key corresponding to this public key.
	PubKey *btcec.PublicKey

	// PrivateTweak is a scalar value that should be added to the private
	// key corresponding to the above public key to obtain the private key
	// to be used to sign this input. This value is typically a leaf node
	// from the revocation tree.
	//
	// NOTE: If this value is nil, then the input can be signed using only
	// the above public key.
	PrivateTweak []byte

	// WitnessScript is the full script required to properly redeem the
	// output. This field will only be populated if a p2wsh or a p2sh
	// output is being signed.
	WitnessScript []byte

	// Output is the target output which should be signed. The PkScript and
	// Value fields within the output should be properly populated,
	// otherwise an invalid signature may be generated.
	Output *wire.TxOut

	// HashType is the target sighash type that should be used when
	// generating the final sighash, and signature.
	HashType txscript.SigHashType

	// SigHashes is the pre-computed sighash midstate to be used when
	// generating the final sighash for signing.
	SigHashes *txscript.TxSigHashes

	// InputIndex is the target input within the transaction that should be
	// signed.
	InputIndex int
}

// Signer represents an abstract object capable of generating raw signatures as
// well as full complete input scripts given a valid SignDescriptor and
// transaction. This interface fully abstracts away signing paving the way for
// Signer implementations such as hardware wallets, hardware tokens, HSM's, or
// simply a regular wallet.
type Signer interface {
	// SignOutputRaw generates a signature for the passed transaction
	// according to the data within the passed SignDescriptor.
	//
	// NOTE: The resulting signature should be void of a sighash byte.
	SignOutputRaw(tx *wire.MsgTx, signDesc *SignDescriptor) ([]byte, error)

	// ComputeInputScript generates a complete InputIndex for the passed
	// transaction with the signature as defined within the passed
	// SignDescriptor. This method should be capable of generating the
	// proper input script for both regular p2wkh output and p2wkh outputs
	// nested within a regular p2sh output.
	ComputeInputScript(tx *wire.MsgTx, signDesc *SignDescriptor) (*InputScript, error)
}

// WalletDriver represents a "driver" for a particular concrete
// WalletController implementation. A driver is identified by a globally unique
// string identifier along with a 'New()' method which is responsible for
// initializing a particular WalletController concrete implementation.
type WalletDriver struct {
	// WalletType is a string which uniquely identifes the WalletController
	// that this driver, drives.
	WalletType string

	// New creates a new instance of a concrete WalletController
	// implementation given a variadic set up arguments. The function takes
	// a varidaic number of interface parameters in order to provide
	// initialization flexibility, thereby accommodating several potential
	// WalletController implementations.
	New func(args ...interface{}) (WalletController, error)
}

var (
	wallets     = make(map[string]*WalletDriver)
	registerMtx sync.Mutex
)

// RegisteredWallets returns a slice of all currently registered notifiers.
//
// NOTE: This function is safe for concurrent access.
func RegisteredWallets() []*WalletDriver {
	registerMtx.Lock()
	defer registerMtx.Unlock()

	registeredWallets := make([]*WalletDriver, 0, len(wallets))
	for _, wallet := range wallets {
		registeredWallets = append(registeredWallets, wallet)
	}

	return registeredWallets
}

// RegisterWallet registers a WalletDriver which is capable of driving a
// concrete WalletController interface. In the case that this driver has
// already been registered, an error is returned.
//
// NOTE: This function is safe for concurrent access.
func RegisterWallet(driver *WalletDriver) error {
	registerMtx.Lock()
	defer registerMtx.Unlock()

	if _, ok := wallets[driver.WalletType]; ok {
		return fmt.Errorf("wallet already registered")
	}

	wallets[driver.WalletType] = driver

	return nil
}

// SupportedWallets returns a slice of strings that represents the wallet
// drivers that have been registered and are therefore supported.
//
// NOTE: This function is safe for concurrent access.
func SupportedWallets() []string {
	registerMtx.Lock()
	defer registerMtx.Unlock()

	supportedWallets := make([]string, 0, len(wallets))
	for walletName := range wallets {
		supportedWallets = append(supportedWallets, walletName)
	}

	return supportedWallets
}
<commit_msg>lnwallet: introduce the MessageSigner interface<commit_after>package lnwallet

import (
	"errors"
	"fmt"
	"sync"

	"github.com/roasbeef/btcd/btcec"
	"github.com/roasbeef/btcd/chaincfg/chainhash"
	"github.com/roasbeef/btcd/txscript"
	"github.com/roasbeef/btcd/wire"
	"github.com/roasbeef/btcutil"
)

// ErrNotMine is an error denoting that a WalletController instance is unable
// to spend a specifid output.
var ErrNotMine = errors.New("the passed output doesn't belong to the wallet")

// AddressType is a enum-like type which denotes the possible address types
// WalletController supports.
type AddressType uint8

const (
	// WitnessPubKey represents a p2wkh address.
	WitnessPubKey AddressType = iota

	// NestedWitnessPubKey represents a p2sh output which is itself a
	// nested p2wkh output.
	NestedWitnessPubKey

	// PubKeyHash represents a regular p2pkh output.
	PubKeyHash
)

// Utxo is an unspent output denoted by its outpoint, and output value of the
// original output.
type Utxo struct {
	Value btcutil.Amount
	wire.OutPoint
}

// TransactionDetail describes a transaction with either inputs which belong to
// the wallet, or has outputs that pay to the wallet.
type TransactionDetail struct {
	// Hash is the transaction hash of the transaction.
	Hash chainhash.Hash

	// Value is the net value of this transaction (in satoshis) from the
	// PoV of the wallet. If this transaction purely spends from the
	// wallet's funds, then this value will be negative. Similarly, if this
	// transaction credits the wallet, then this value will be positive.
	Value btcutil.Amount

	// NumConfirmations is the number of confirmations this transaction
	// has. If the transaction is unconfirmed, then this value will be
	// zero.
	NumConfirmations int32

	// BlockHeight is the hash of the block which includes this
	// transaction. Unconfirmed transactions will have a nil value for this
	// field.
	BlockHash *chainhash.Hash

	// BlockHeight is the height of the block including this transaction.
	// Unconfirmed transaction will show a height of zero.
	BlockHeight int32

	// Timestamp is the unix timestamp of the block including this
	// transaction. If the transaction is unconfirmed, then this will be a
	// timestamp of txn creation.
	Timestamp int64

	// TotalFees is the total fee in satoshis paid by this transaction.
	TotalFees int64
}

// TransactionSubscription is an interface which describes an object capable of
// receiving notifications of new transaction related to the underlying wallet.
// TODO(roasbeef): add balance updates?
type TransactionSubscription interface {
	// ConfirmedTransactions returns a channel which will be sent on as new
	// relevant transactions are confirmed.
	ConfirmedTransactions() chan *TransactionDetail

	// UnconfirmedTransactions returns a channel which will be sent on as
	// new relevant transactions are seen within the network.
	UnconfirmedTransactions() chan *TransactionDetail

	// Cancel finalizes the subscription, cleaning up any resources
	// allocated.
	Cancel()
}

// WalletController defines an abstract interface for controlling a local Pure
// Go wallet, a local or remote wallet via an RPC mechanism, or possibly even
// a daemon assisted hardware wallet. This interface serves the purpose of
// allowing LightningWallet to be seamlessly compatible with several wallets
// such as: uspv, btcwallet, Bitcoin Core, Electrum, etc. This interface then
// serves as a "base wallet", with Lightning Network awareness taking place at
// a "higher" level of abstraction. Essentially, an overlay wallet.
// Implementors of this interface must closely adhere to the documented
// behavior of all interface methods in order to ensure identical behavior
// across all concrete implementations.
type WalletController interface {
	// FetchInputInfo queries for the WalletController's knowledge of the
	// passed outpoint. If the base wallet determines this output is under
	// its control, then the original txout should be returned. Otherwise,
	// a non-nil error value of ErrNotMine should be returned instead.
	FetchInputInfo(prevOut *wire.OutPoint) (*wire.TxOut, error)

	// ConfirmedBalance returns the sum of all the wallet's unspent outputs
	// that have at least confs confirmations. If confs is set to zero,
	// then all unspent outputs, including those currently in the mempool
	// will be included in the final sum.
	ConfirmedBalance(confs int32, witness bool) (btcutil.Amount, error)

	// NewAddress returns the next external or internal address for the
	// wallet dictated by the value of the `change` parameter. If change is
	// true, then an internal address should be used, otherwise an external
	// address should be returned. The type of address returned is dictated
	// by the wallet's capabilities, and may be of type: p2sh, p2pkh,
	// p2wkh, p2wsh, etc.
	NewAddress(addrType AddressType, change bool) (btcutil.Address, error)

	// GetPrivKey retrives the underlying private key associated with the
	// passed address. If the wallet is unable to locate this private key
	// due to the address not being under control of the wallet, then an
	// error should be returned.
	GetPrivKey(a btcutil.Address) (*btcec.PrivateKey, error)

	// NewRawKey returns a raw private key controlled by the wallet. These
	// keys are used for the 2-of-2 multi-sig outputs for funding
	// transactions, as well as the pub key used for commitment transactions.
	//
	// NOTE: The wallet MUST watch for on-chain outputs created to a p2wpkh
	// script using keys returned by this function.
	NewRawKey() (*btcec.PublicKey, error)

	// FetchRootKey returns a root key which will be used by the
	// LightningWallet to deterministically generate secrets. The private
	// key returned by this method should remain constant in-between
	// WalletController restarts.
	FetchRootKey() (*btcec.PrivateKey, error)

	// SendOutputs funds, signs, and broadcasts a Bitcoin transaction
	// paying out to the specified outputs. In the case the wallet has
	// insufficient funds, or the outputs are non-standard, an error
	// should be returned.
	SendOutputs(outputs []*wire.TxOut) (*chainhash.Hash, error)

	// ListUnspentWitness returns all unspent outputs which are version 0
	// witness programs. The 'confirms' parameter indicates the minimum
	// number of confirmations an output needs in order to be returned by
	// this method. Passing -1 as 'confirms' indicates that even
	// unconfirmed outputs should be returned.
	ListUnspentWitness(confirms int32) ([]*Utxo, error)

	// ListTransactionDetails returns a list of all transactions which are
	// relevant to the wallet.
	ListTransactionDetails() ([]*TransactionDetail, error)

	// LockOutpoint marks an outpoint as locked meaning it will no longer
	// be deemed as eligible for coin selection. Locking outputs are
	// utilized in order to avoid race conditions when selecting inputs for
	// usage when funding a channel.
	LockOutpoint(o wire.OutPoint)

	// UnlockOutpoint unlocks an previously locked output, marking it
	// eligible for coin selection.
	UnlockOutpoint(o wire.OutPoint)

	// PublishTransaction performs cursory validation (dust checks, etc),
	// then finally broadcasts the passed transaction to the Bitcoin network.
	PublishTransaction(tx *wire.MsgTx) error

	// SubscribeTransactions returns a TransactionSubscription client which
	// is capable of receiving async notifications as new transactions
	// related to the wallet are seen within the network, or found in
	// blocks.
	//
	// NOTE: a non-nil error should be returned if notifications aren't
	// supported.
	//
	// TODO(roasbeef): make distinct interface?
	SubscribeTransactions() (TransactionSubscription, error)

	// IsSynced returns a boolean indicating if from the PoV of the wallet,
	// it has fully synced to the current best block in the main chain.
	IsSynced() (bool, error)

	// Start initializes the wallet, making any necessary connections,
	// starting up required goroutines etc.
	Start() error

	// Stop signals the wallet for shutdown. Shutdown may entail closing
	// any active sockets, database handles, stopping goroutines, etc.
	Stop() error
}

// BlockChainIO is a dedicated source which will be used to obtain queries
// related to the current state of the blockchain. The data returned by each of
// the defined methods within this interface should always return the most up
// to date data possible.
//
// TODO(roasbeef): move to diff package perhaps?
// TODO(roasbeef): move publish txn here?
type BlockChainIO interface {
	// GetBestBlock returns the current height and block hash of the valid
	// most-work chain the implementation is aware of.
	GetBestBlock() (*chainhash.Hash, int32, error)

	// GetTxOut returns the original output referenced by the passed
	// outpoint.
	GetUtxo(txid *chainhash.Hash, index uint32) (*wire.TxOut, error)

	// GetTransaction returns the full transaction identified by the passed
	// transaction ID.
	GetTransaction(txid *chainhash.Hash) (*wire.MsgTx, error)

	// GetBlockHash returns the hash of the block in the best blockchain
	// at the given height.
	GetBlockHash(blockHeight int64) (*chainhash.Hash, error)

	// GetBlock returns the block in the main chain identified by the given
	// hash.
	GetBlock(blockHash *chainhash.Hash) (*wire.MsgBlock, error)
}

// SignDescriptor houses the necessary information required to successfully sign
// a given output. This struct is used by the Signer interface in order to gain
// access to critical data needed to generate a valid signature.
type SignDescriptor struct {
	// Pubkey is the public key to which the signature should be generated
	// over. The Signer should then generate a signature with the private
	// key corresponding to this public key.
	PubKey *btcec.PublicKey

	// PrivateTweak is a scalar value that should be added to the private
	// key corresponding to the above public key to obtain the private key
	// to be used to sign this input. This value is typically a leaf node
	// from the revocation tree.
	//
	// NOTE: If this value is nil, then the input can be signed using only
	// the above public key.
	PrivateTweak []byte

	// WitnessScript is the full script required to properly redeem the
	// output. This field will only be populated if a p2wsh or a p2sh
	// output is being signed.
	WitnessScript []byte

	// Output is the target output which should be signed. The PkScript and
	// Value fields within the output should be properly populated,
	// otherwise an invalid signature may be generated.
	Output *wire.TxOut

	// HashType is the target sighash type that should be used when
	// generating the final sighash, and signature.
	HashType txscript.SigHashType

	// SigHashes is the pre-computed sighash midstate to be used when
	// generating the final sighash for signing.
	SigHashes *txscript.TxSigHashes

	// InputIndex is the target input within the transaction that should be
	// signed.
	InputIndex int
}

// Signer represents an abstract object capable of generating raw signatures as
// well as full complete input scripts given a valid SignDescriptor and
// transaction. This interface fully abstracts away signing paving the way for
// Signer implementations such as hardware wallets, hardware tokens, HSM's, or
// simply a regular wallet.
type Signer interface {
	// SignOutputRaw generates a signature for the passed transaction
	// according to the data within the passed SignDescriptor.
	//
	// NOTE: The resulting signature should be void of a sighash byte.
	SignOutputRaw(tx *wire.MsgTx, signDesc *SignDescriptor) ([]byte, error)

	// ComputeInputScript generates a complete InputIndex for the passed
	// transaction with the signature as defined within the passed
	// SignDescriptor. This method should be capable of generating the
	// proper input script for both regular p2wkh output and p2wkh outputs
	// nested within a regular p2sh output.
	ComputeInputScript(tx *wire.MsgTx, signDesc *SignDescriptor) (*InputScript, error)
}

// MessageSigner represents an abstract object capable of signing arbitrary
// messages. The capabilities of this interface are used to sign announcements
// to the network, or just arbitrary messages that leverage the wallet's keys
// to attest to some message.
type MessageSigner interface {
	// SignMessage attempts to sign a target message with the private key
	// that corresponds to the passed public key. If the target private key
	// is unable to be found, then an error will be returned. The actual
	// digest signed is the double SHA-256 of the passed message.
	SignMessage(pubKey *btcec.PublicKey, msg []byte) (*btcec.Signature, error)
}

// WalletDriver represents a "driver" for a particular concrete
// WalletController implementation. A driver is identified by a globally unique
// string identifier along with a 'New()' method which is responsible for
// initializing a particular WalletController concrete implementation.
type WalletDriver struct {
	// WalletType is a string which uniquely identifes the WalletController
	// that this driver, drives.
	WalletType string

	// New creates a new instance of a concrete WalletController
	// implementation given a variadic set up arguments. The function takes
	// a varidaic number of interface parameters in order to provide
	// initialization flexibility, thereby accommodating several potential
	// WalletController implementations.
	New func(args ...interface{}) (WalletController, error)
}

var (
	wallets     = make(map[string]*WalletDriver)
	registerMtx sync.Mutex
)

// RegisteredWallets returns a slice of all currently registered notifiers.
//
// NOTE: This function is safe for concurrent access.
func RegisteredWallets() []*WalletDriver {
	registerMtx.Lock()
	defer registerMtx.Unlock()

	registeredWallets := make([]*WalletDriver, 0, len(wallets))
	for _, wallet := range wallets {
		registeredWallets = append(registeredWallets, wallet)
	}

	return registeredWallets
}

// RegisterWallet registers a WalletDriver which is capable of driving a
// concrete WalletController interface. In the case that this driver has
// already been registered, an error is returned.
//
// NOTE: This function is safe for concurrent access.
func RegisterWallet(driver *WalletDriver) error {
	registerMtx.Lock()
	defer registerMtx.Unlock()

	if _, ok := wallets[driver.WalletType]; ok {
		return fmt.Errorf("wallet already registered")
	}

	wallets[driver.WalletType] = driver

	return nil
}

// SupportedWallets returns a slice of strings that represents the wallet
// drivers that have been registered and are therefore supported.
//
// NOTE: This function is safe for concurrent access.
func SupportedWallets() []string {
	registerMtx.Lock()
	defer registerMtx.Unlock()

	supportedWallets := make([]string, 0, len(wallets))
	for walletName := range wallets {
		supportedWallets = append(supportedWallets, walletName)
	}

	return supportedWallets
}
<|endoftext|>


<commit_before>package cluster

import (
	"context"
	"sync"
	"time"

	"github.com/lxc/lxd/client"
	"github.com/lxc/lxd/lxd/db"
	"github.com/lxc/lxd/lxd/endpoints"
	"github.com/lxc/lxd/lxd/task"
	"github.com/lxc/lxd/shared"
	"github.com/lxc/lxd/shared/api"
	log "github.com/lxc/lxd/shared/log15"
	"github.com/lxc/lxd/shared/logger"
)

var listeners = map[string]*lxd.EventListener{}
var listenersLock sync.Mutex

// Events starts a task that continuously monitors the list of cluster nodes and
// maintains a pool of websocket connections against all of them, in order to
// get notified about events.
//
// Whenever an event is received the given callback is invoked.
func Events(endpoints *endpoints.Endpoints, cluster *db.Cluster, serverCert func() *shared.CertInfo, f func(int64, api.Event)) (task.Func, task.Schedule) {
	// Update our pool of event listeners. Since database queries are
	// blocking, we spawn the actual logic in a goroutine, to abort
	// immediately when we receive the stop signal.
	update := func(ctx context.Context) {
		ch := make(chan struct{})
		go func() {
			eventsUpdateListeners(endpoints, cluster, serverCert, nil, f)
			ch <- struct{}{}
		}()
		select {
		case <-ch:
		case <-ctx.Done():
		}
	}

	schedule := task.Every(time.Second)

	return update, schedule
}

func eventsUpdateListeners(endpoints *endpoints.Endpoints, cluster *db.Cluster, serverCert func() *shared.CertInfo, members map[int64]APIHeartbeatMember, f func(int64, api.Event)) {
	// Get the current cluster nodes.
	var nodes []db.NodeInfo
	var offlineThreshold time.Duration

	err := cluster.Transaction(func(tx *db.ClusterTx) error {
		var err error

		nodes, err = tx.GetNodes()
		if err != nil {
			return err
		}

		offlineThreshold, err = tx.GetNodeOfflineThreshold()
		if err != nil {
			return err
		}

		return nil
	})
	if err != nil {
		logger.Warn("Failed to get current cluster members", log.Ctx{"err": err})
		return
	}
	if len(nodes) == 1 {
		return // Either we're not clustered or this is a single-node cluster
	}

	address := endpoints.NetworkAddress()

	addresses := make([]string, len(nodes))
	for i, node := range nodes {
		addresses[i] = node.Address

		if node.Address == address {
			continue
		}

		listenersLock.Lock()
		listener, ok := listeners[node.Address]

		// Don't bother trying to connect to offline nodes, or to ourselves.
		if node.IsOffline(offlineThreshold) {
			if ok {
				listener.Disconnect()
			}

			listenersLock.Unlock()
			continue
		}

		// The node has already a listener associated to it.
		if ok {
			// Double check that the listener is still
			// connected. If it is, just move on, other
			// we'll try to connect again.
			if listeners[node.Address].IsActive() {
				listenersLock.Unlock()
				continue
			}

			delete(listeners, node.Address)
		}
		listenersLock.Unlock()

		listener, err := eventsConnect(node.Address, endpoints.NetworkCert(), serverCert())
		if err != nil {
			logger.Warn("Failed to get events from member", log.Ctx{"address": node.Address, "err": err})
			continue
		}
		logger.Debug("Listening for events on member", log.Ctx{"address": node.Address})
		listener.AddHandler(nil, func(event api.Event) { f(node.ID, event) })

		listenersLock.Lock()
		listeners[node.Address] = listener
		listenersLock.Unlock()
	}

	listenersLock.Lock()
	for address, listener := range listeners {
		if !shared.StringInSlice(address, addresses) {
			listener.Disconnect()
			delete(listeners, address)
		}
	}
	listenersLock.Unlock()
}

// Establish a client connection to get events from the given node.
func eventsConnect(address string, networkCert *shared.CertInfo, serverCert *shared.CertInfo) (*lxd.EventListener, error) {
	client, err := Connect(address, networkCert, serverCert, nil, true)
	if err != nil {
		return nil, err
	}

	// Set the project to the special wildcard in order to get notified
	// about all events across all projects.
	client = client.UseProject("*")

	return client.GetEvents()
}
<commit_msg>lxd/cluster/events: Load members from global DB if no heartbeat members provided in eventsUpdateListeners<commit_after>package cluster

import (
	"context"
	"sync"
	"time"

	"github.com/lxc/lxd/client"
	"github.com/lxc/lxd/lxd/db"
	"github.com/lxc/lxd/lxd/endpoints"
	"github.com/lxc/lxd/lxd/task"
	"github.com/lxc/lxd/shared"
	"github.com/lxc/lxd/shared/api"
	log "github.com/lxc/lxd/shared/log15"
	"github.com/lxc/lxd/shared/logger"
)

var listeners = map[string]*lxd.EventListener{}
var listenersLock sync.Mutex

// Events starts a task that continuously monitors the list of cluster nodes and
// maintains a pool of websocket connections against all of them, in order to
// get notified about events.
//
// Whenever an event is received the given callback is invoked.
func Events(endpoints *endpoints.Endpoints, cluster *db.Cluster, serverCert func() *shared.CertInfo, f func(int64, api.Event)) (task.Func, task.Schedule) {
	// Update our pool of event listeners. Since database queries are
	// blocking, we spawn the actual logic in a goroutine, to abort
	// immediately when we receive the stop signal.
	update := func(ctx context.Context) {
		ch := make(chan struct{})
		go func() {
			eventsUpdateListeners(endpoints, cluster, serverCert, nil, f)
			ch <- struct{}{}
		}()
		select {
		case <-ch:
		case <-ctx.Done():
		}
	}

	schedule := task.Every(time.Second)

	return update, schedule
}

func eventsUpdateListeners(endpoints *endpoints.Endpoints, cluster *db.Cluster, serverCert func() *shared.CertInfo, members map[int64]APIHeartbeatMember, f func(int64, api.Event)) {
	// If no heartbeat members provided, populate from global database.
	if members == nil {
		var dbMembers []db.NodeInfo
		var offlineThreshold time.Duration

		err := cluster.Transaction(func(tx *db.ClusterTx) error {
			var err error

			dbMembers, err = tx.GetNodes()
			if err != nil {
				return err
			}

			offlineThreshold, err = tx.GetNodeOfflineThreshold()
			if err != nil {
				return err
			}

			return nil
		})
		if err != nil {
			logger.Warn("Failed to get current cluster members", log.Ctx{"err": err})
			return
		}

		members = make(map[int64]APIHeartbeatMember, len(dbMembers))
		for _, dbMember := range dbMembers {
			members[dbMember.ID] = APIHeartbeatMember{
				ID:            dbMember.ID,
				Name:          dbMember.Name,
				Address:       dbMember.Address,
				LastHeartbeat: dbMember.Heartbeat,
				Online:        !dbMember.IsOffline(offlineThreshold),
			}
		}
	}

	// Get the current cluster nodes.
	var nodes []db.NodeInfo
	var offlineThreshold time.Duration

	err := cluster.Transaction(func(tx *db.ClusterTx) error {
		var err error

		nodes, err = tx.GetNodes()
		if err != nil {
			return err
		}

		offlineThreshold, err = tx.GetNodeOfflineThreshold()
		if err != nil {
			return err
		}

		return nil
	})
	if err != nil {
		logger.Warn("Failed to get current cluster members", log.Ctx{"err": err})
		return
	}
	if len(nodes) == 1 {
		return // Either we're not clustered or this is a single-node cluster
	}

	address := endpoints.NetworkAddress()

	addresses := make([]string, len(nodes))
	for i, node := range nodes {
		addresses[i] = node.Address

		if node.Address == address {
			continue
		}

		listenersLock.Lock()
		listener, ok := listeners[node.Address]

		// Don't bother trying to connect to offline nodes, or to ourselves.
		if node.IsOffline(offlineThreshold) {
			if ok {
				listener.Disconnect()
			}

			listenersLock.Unlock()
			continue
		}

		// The node has already a listener associated to it.
		if ok {
			// Double check that the listener is still
			// connected. If it is, just move on, other
			// we'll try to connect again.
			if listeners[node.Address].IsActive() {
				listenersLock.Unlock()
				continue
			}

			delete(listeners, node.Address)
		}
		listenersLock.Unlock()

		listener, err := eventsConnect(node.Address, endpoints.NetworkCert(), serverCert())
		if err != nil {
			logger.Warn("Failed to get events from member", log.Ctx{"address": node.Address, "err": err})
			continue
		}
		logger.Debug("Listening for events on member", log.Ctx{"address": node.Address})
		listener.AddHandler(nil, func(event api.Event) { f(node.ID, event) })

		listenersLock.Lock()
		listeners[node.Address] = listener
		listenersLock.Unlock()
	}

	listenersLock.Lock()
	for address, listener := range listeners {
		if !shared.StringInSlice(address, addresses) {
			listener.Disconnect()
			delete(listeners, address)
		}
	}
	listenersLock.Unlock()
}

// Establish a client connection to get events from the given node.
func eventsConnect(address string, networkCert *shared.CertInfo, serverCert *shared.CertInfo) (*lxd.EventListener, error) {
	client, err := Connect(address, networkCert, serverCert, nil, true)
	if err != nil {
		return nil, err
	}

	// Set the project to the special wildcard in order to get notified
	// about all events across all projects.
	client = client.UseProject("*")

	return client.GetEvents()
}
<|endoftext|>


<commit_before>package device

import (
	"fmt"

	deviceConfig "github.com/lxc/lxd/lxd/device/config"
	"github.com/lxc/lxd/lxd/instance"
	"github.com/lxc/lxd/lxd/instance/instancetype"
	"github.com/lxc/lxd/lxd/network"
	"github.com/lxc/lxd/lxd/revert"
	"github.com/lxc/lxd/lxd/util"
)

type nicP2P struct {
	deviceCommon
}

// CanHotPlug returns whether the device can be managed whilst the instance is running. Returns true.
func (d *nicP2P) CanHotPlug() bool {
	return true
}

// validateConfig checks the supplied config for correctness.
func (d *nicP2P) validateConfig(instConf instance.ConfigReader) error {
	if !instanceSupported(instConf.Type(), instancetype.Container, instancetype.VM) {
		return ErrUnsupportedDevType
	}

	optionalFields := []string{
		"name",
		"mtu",
		"hwaddr",
		"host_name",
		"limits.ingress",
		"limits.egress",
		"limits.max",
		"ipv4.routes",
		"ipv6.routes",
		"boot.priority",
	}
	err := d.config.Validate(nicValidationRules([]string{}, optionalFields, instConf))
	if err != nil {
		return err
	}

	return nil
}

// validateEnvironment checks the runtime environment for correctness.
func (d *nicP2P) validateEnvironment() error {
	if d.inst.Type() == instancetype.Container && d.config["name"] == "" {
		return fmt.Errorf("Requires name property to start")
	}

	return nil
}

// UpdatableFields returns a list of fields that can be updated without triggering a device remove & add.
func (d *nicP2P) UpdatableFields(oldDevice Type) []string {
	// Check old and new device types match.
	_, match := oldDevice.(*nicP2P)
	if !match {
		return []string{}
	}

	return []string{"limits.ingress", "limits.egress", "limits.max", "ipv4.routes", "ipv6.routes"}
}

// Start is run when the device is added to a running instance or instance is starting up.
func (d *nicP2P) Start() (*deviceConfig.RunConfig, error) {
	err := d.validateEnvironment()
	if err != nil {
		return nil, err
	}

	revert := revert.New()
	defer revert.Fail()

	saveData := make(map[string]string)
	saveData["host_name"] = d.config["host_name"]

	var peerName string

	// Create veth pair and configure the peer end with custom hwaddr and mtu if supplied.
	if d.inst.Type() == instancetype.Container {
		if saveData["host_name"] == "" {
			saveData["host_name"] = network.RandomDevName("veth")
		}
		peerName, err = networkCreateVethPair(saveData["host_name"], d.config)
	} else if d.inst.Type() == instancetype.VM {
		if saveData["host_name"] == "" {
			saveData["host_name"] = network.RandomDevName("tap")
		}
		peerName = saveData["host_name"] // VMs use the host_name to link to the TAP FD.
		err = networkCreateTap(saveData["host_name"], d.config)
	}

	if err != nil {
		return nil, err
	}

	revert.Add(func() { network.InterfaceRemove(saveData["host_name"]) })

	// Populate device config with volatile fields if needed.
	networkVethFillFromVolatile(d.config, saveData)

	// Apply host-side routes to veth interface.
	err = networkNICRouteAdd(d.config["host_name"], append(util.SplitNTrimSpace(d.config["ipv4.routes"], ",", -1, true), util.SplitNTrimSpace(d.config["ipv6.routes"], ",", -1, true)...)...)
	if err != nil {
		return nil, err
	}

	// Apply host-side limits.
	err = networkSetupHostVethLimits(d.config)
	if err != nil {
		return nil, err
	}

	err = d.volatileSet(saveData)
	if err != nil {
		return nil, err
	}

	runConf := deviceConfig.RunConfig{}
	runConf.NetworkInterface = []deviceConfig.RunConfigItem{
		{Key: "type", Value: "phys"},
		{Key: "name", Value: d.config["name"]},
		{Key: "flags", Value: "up"},
		{Key: "link", Value: peerName},
	}

	if d.inst.Type() == instancetype.VM {
		runConf.NetworkInterface = append(runConf.NetworkInterface,
			[]deviceConfig.RunConfigItem{
				{Key: "devName", Value: d.name},
				{Key: "hwaddr", Value: d.config["hwaddr"]},
			}...)
	}

	revert.Success()
	return &runConf, nil
}

// Update applies configuration changes to a started device.
func (d *nicP2P) Update(oldDevices deviceConfig.Devices, isRunning bool) error {
	if !isRunning {
		return nil
	}

	err := d.validateEnvironment()
	if err != nil {
		return err
	}

	oldConfig := oldDevices[d.name]
	v := d.volatileGet()

	// Populate device config with volatile fields if needed.
	networkVethFillFromVolatile(d.config, v)
	networkVethFillFromVolatile(oldConfig, v)

	// Remove old host-side routes from veth interface.
	networkNICRouteDelete(oldConfig["host_name"], append(util.SplitNTrimSpace(oldConfig["ipv4.routes"], ",", -1, true), util.SplitNTrimSpace(oldConfig["ipv6.routes"], ",", -1, true)...)...)

	// Apply host-side routes to veth interface.
	err = networkNICRouteAdd(d.config["host_name"], append(util.SplitNTrimSpace(d.config["ipv4.routes"], ",", -1, true), util.SplitNTrimSpace(d.config["ipv6.routes"], ",", -1, true)...)...)
	if err != nil {
		return err
	}

	// Apply host-side limits.
	err = networkSetupHostVethLimits(d.config)
	if err != nil {
		return err
	}

	return nil
}

// Stop is run when the device is removed from the instance.
func (d *nicP2P) Stop() (*deviceConfig.RunConfig, error) {
	runConf := deviceConfig.RunConfig{
		PostHooks: []func() error{d.postStop},
	}

	return &runConf, nil
}

// postStop is run after the device is removed from the instance.
func (d *nicP2P) postStop() error {
	defer d.volatileSet(map[string]string{
		"host_name": "",
	})

	v := d.volatileGet()

	networkVethFillFromVolatile(d.config, v)

	if d.config["host_name"] != "" && network.InterfaceExists(d.config["host_name"]) {
		// Removing host-side end of veth pair will delete the peer end too.
		err := network.InterfaceRemove(d.config["host_name"])
		if err != nil {
			return fmt.Errorf("Failed to remove interface %s: %s", d.config["host_name"], err)
		}
	}

	return nil
}
<commit_msg>lxd/device/nic/p2p: networkCreateVethPair and networkCreateTap usage<commit_after>package device

import (
	"fmt"

	deviceConfig "github.com/lxc/lxd/lxd/device/config"
	"github.com/lxc/lxd/lxd/instance"
	"github.com/lxc/lxd/lxd/instance/instancetype"
	"github.com/lxc/lxd/lxd/network"
	"github.com/lxc/lxd/lxd/revert"
	"github.com/lxc/lxd/lxd/util"
)

type nicP2P struct {
	deviceCommon
}

// CanHotPlug returns whether the device can be managed whilst the instance is running. Returns true.
func (d *nicP2P) CanHotPlug() bool {
	return true
}

// validateConfig checks the supplied config for correctness.
func (d *nicP2P) validateConfig(instConf instance.ConfigReader) error {
	if !instanceSupported(instConf.Type(), instancetype.Container, instancetype.VM) {
		return ErrUnsupportedDevType
	}

	optionalFields := []string{
		"name",
		"mtu",
		"hwaddr",
		"host_name",
		"limits.ingress",
		"limits.egress",
		"limits.max",
		"ipv4.routes",
		"ipv6.routes",
		"boot.priority",
	}
	err := d.config.Validate(nicValidationRules([]string{}, optionalFields, instConf))
	if err != nil {
		return err
	}

	return nil
}

// validateEnvironment checks the runtime environment for correctness.
func (d *nicP2P) validateEnvironment() error {
	if d.inst.Type() == instancetype.Container && d.config["name"] == "" {
		return fmt.Errorf("Requires name property to start")
	}

	return nil
}

// UpdatableFields returns a list of fields that can be updated without triggering a device remove & add.
func (d *nicP2P) UpdatableFields(oldDevice Type) []string {
	// Check old and new device types match.
	_, match := oldDevice.(*nicP2P)
	if !match {
		return []string{}
	}

	return []string{"limits.ingress", "limits.egress", "limits.max", "ipv4.routes", "ipv6.routes"}
}

// Start is run when the device is added to a running instance or instance is starting up.
func (d *nicP2P) Start() (*deviceConfig.RunConfig, error) {
	err := d.validateEnvironment()
	if err != nil {
		return nil, err
	}

	revert := revert.New()
	defer revert.Fail()

	saveData := make(map[string]string)
	saveData["host_name"] = d.config["host_name"]

	var peerName string
	var mtu uint32

	// Create veth pair and configure the peer end with custom hwaddr and mtu if supplied.
	if d.inst.Type() == instancetype.Container {
		if saveData["host_name"] == "" {
			saveData["host_name"] = network.RandomDevName("veth")
		}

		peerName, mtu, err = networkCreateVethPair(saveData["host_name"], d.config)
	} else if d.inst.Type() == instancetype.VM {
		if saveData["host_name"] == "" {
			saveData["host_name"] = network.RandomDevName("tap")
		}

		peerName = saveData["host_name"] // VMs use the host_name to link to the TAP FD.
		mtu, err = networkCreateTap(saveData["host_name"], d.config)
	}

	if err != nil {
		return nil, err
	}

	revert.Add(func() { network.InterfaceRemove(saveData["host_name"]) })

	// Populate device config with volatile fields if needed.
	networkVethFillFromVolatile(d.config, saveData)

	// Apply host-side routes to veth interface.
	err = networkNICRouteAdd(d.config["host_name"], append(util.SplitNTrimSpace(d.config["ipv4.routes"], ",", -1, true), util.SplitNTrimSpace(d.config["ipv6.routes"], ",", -1, true)...)...)
	if err != nil {
		return nil, err
	}

	// Apply host-side limits.
	err = networkSetupHostVethLimits(d.config)
	if err != nil {
		return nil, err
	}

	err = d.volatileSet(saveData)
	if err != nil {
		return nil, err
	}

	runConf := deviceConfig.RunConfig{}
	runConf.NetworkInterface = []deviceConfig.RunConfigItem{
		{Key: "type", Value: "phys"},
		{Key: "name", Value: d.config["name"]},
		{Key: "flags", Value: "up"},
		{Key: "link", Value: peerName},
	}

	if d.inst.Type() == instancetype.VM {
		runConf.NetworkInterface = append(runConf.NetworkInterface,
			[]deviceConfig.RunConfigItem{
				{Key: "devName", Value: d.name},
				{Key: "hwaddr", Value: d.config["hwaddr"]},
				{Key: "mtu", Value: fmt.Sprintf("%d", mtu)},
			}...)
	}

	revert.Success()
	return &runConf, nil
}

// Update applies configuration changes to a started device.
func (d *nicP2P) Update(oldDevices deviceConfig.Devices, isRunning bool) error {
	if !isRunning {
		return nil
	}

	err := d.validateEnvironment()
	if err != nil {
		return err
	}

	oldConfig := oldDevices[d.name]
	v := d.volatileGet()

	// Populate device config with volatile fields if needed.
	networkVethFillFromVolatile(d.config, v)
	networkVethFillFromVolatile(oldConfig, v)

	// Remove old host-side routes from veth interface.
	networkNICRouteDelete(oldConfig["host_name"], append(util.SplitNTrimSpace(oldConfig["ipv4.routes"], ",", -1, true), util.SplitNTrimSpace(oldConfig["ipv6.routes"], ",", -1, true)...)...)

	// Apply host-side routes to veth interface.
	err = networkNICRouteAdd(d.config["host_name"], append(util.SplitNTrimSpace(d.config["ipv4.routes"], ",", -1, true), util.SplitNTrimSpace(d.config["ipv6.routes"], ",", -1, true)...)...)
	if err != nil {
		return err
	}

	// Apply host-side limits.
	err = networkSetupHostVethLimits(d.config)
	if err != nil {
		return err
	}

	return nil
}

// Stop is run when the device is removed from the instance.
func (d *nicP2P) Stop() (*deviceConfig.RunConfig, error) {
	runConf := deviceConfig.RunConfig{
		PostHooks: []func() error{d.postStop},
	}

	return &runConf, nil
}

// postStop is run after the device is removed from the instance.
func (d *nicP2P) postStop() error {
	defer d.volatileSet(map[string]string{
		"host_name": "",
	})

	v := d.volatileGet()

	networkVethFillFromVolatile(d.config, v)

	if d.config["host_name"] != "" && network.InterfaceExists(d.config["host_name"]) {
		// Removing host-side end of veth pair will delete the peer end too.
		err := network.InterfaceRemove(d.config["host_name"])
		if err != nil {
			return fmt.Errorf("Failed to remove interface %s: %s", d.config["host_name"], err)
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package amazonec2

import (
	"testing"

	"errors"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/docker/machine/commands/commandstest"
	"github.com/stretchr/testify/assert"
)

const (
	testSSHPort    = 22
	testDockerPort = 2376
)

var (
	securityGroup = &ec2.SecurityGroup{
		GroupName: aws.String("test-group"),
		GroupId:   aws.String("12345"),
		VpcId:     aws.String("12345"),
	}
)

func TestConfigureSecurityGroupPermissionsEmpty(t *testing.T) {
	driver := NewTestDriver()

	perms := driver.configureSecurityGroupPermissions(securityGroup)

	assert.Len(t, perms, 2)
}

func TestConfigureSecurityGroupPermissionsSshOnly(t *testing.T) {
	driver := NewTestDriver()
	group := securityGroup
	group.IpPermissions = []*ec2.IpPermission{
		{
			IpProtocol: aws.String("tcp"),
			FromPort:   aws.Int64(int64(testSSHPort)),
			ToPort:     aws.Int64(int64(testSSHPort)),
		},
	}

	perms := driver.configureSecurityGroupPermissions(group)

	assert.Len(t, perms, 1)
	assert.Equal(t, testDockerPort, *perms[0].FromPort)
}

func TestConfigureSecurityGroupPermissionsDockerOnly(t *testing.T) {
	driver := NewTestDriver()
	group := securityGroup
	group.IpPermissions = []*ec2.IpPermission{
		{
			IpProtocol: aws.String("tcp"),
			FromPort:   aws.Int64((testDockerPort)),
			ToPort:     aws.Int64((testDockerPort)),
		},
	}

	perms := driver.configureSecurityGroupPermissions(group)

	assert.Len(t, perms, 1)
	assert.Equal(t, testSSHPort, *perms[0].FromPort)
}

func TestConfigureSecurityGroupPermissionsDockerAndSsh(t *testing.T) {
	driver := NewTestDriver()
	group := securityGroup
	group.IpPermissions = []*ec2.IpPermission{
		{
			IpProtocol: aws.String("tcp"),
			FromPort:   aws.Int64(testSSHPort),
			ToPort:     aws.Int64(testSSHPort),
		},
		{
			IpProtocol: aws.String("tcp"),
			FromPort:   aws.Int64(testDockerPort),
			ToPort:     aws.Int64(testDockerPort),
		},
	}

	perms := driver.configureSecurityGroupPermissions(group)

	assert.Empty(t, perms)
}

func TestValidateAwsRegionValid(t *testing.T) {
	regions := []string{"eu-west-1", "eu-central-1"}

	for _, region := range regions {
		validatedRegion, err := validateAwsRegion(region)

		assert.NoError(t, err)
		assert.Equal(t, region, validatedRegion)
	}
}

func TestValidateAwsRegionInvalid(t *testing.T) {
	regions := []string{"eu-west-2", "eu-central-2"}

	for _, region := range regions {
		_, err := validateAwsRegion(region)

		assert.EqualError(t, err, "Invalid region specified")
	}
}

func TestFindDefaultVPC(t *testing.T) {
	driver := NewDriver("machineFoo", "path")
	driver.clientFactory = func() Ec2Client { return &fakeEC2WithLogin{} }

	vpc, err := driver.getDefaultVPCId()

	assert.Equal(t, "vpc-9999", vpc)
	assert.NoError(t, err)
}

func TestDefaultVPCIsMissing(t *testing.T) {
	driver := NewDriver("machineFoo", "path")
	driver.clientFactory = func() Ec2Client {
		return &fakeEC2WithDescribe{
			output: &ec2.DescribeAccountAttributesOutput{
				AccountAttributes: []*ec2.AccountAttribute{},
			},
		}
	}

	vpc, err := driver.getDefaultVPCId()

	assert.EqualError(t, err, "No default-vpc attribute")
	assert.Empty(t, vpc)
}

func TestDescribeAccountAttributeFails(t *testing.T) {
	driver := NewDriver("machineFoo", "path")
	driver.clientFactory = func() Ec2Client {
		return &fakeEC2WithDescribe{
			err: errors.New("Not Found"),
		}
	}

	vpc, err := driver.getDefaultVPCId()

	assert.EqualError(t, err, "Not Found")
	assert.Empty(t, vpc)
}

func TestAccessKeyIsMandatory(t *testing.T) {
	driver := NewTestDriver()
	driver.awsCredentials = &cliCredentials{}
	options := &commandstest.FakeFlagger{
		Data: map[string]interface{}{
			"name":             "test",
			"amazonec2-region": "us-east-1",
			"amazonec2-zone":   "e",
		},
	}

	err := driver.SetConfigFromFlags(options)

	assert.Equal(t, err, errorMissingAccessKeyOption)
}

func TestAccessKeyIsMandatoryEvenIfSecretKeyIsPassed(t *testing.T) {
	driver := NewTestDriver()
	driver.awsCredentials = &cliCredentials{}
	options := &commandstest.FakeFlagger{
		Data: map[string]interface{}{
			"name":                 "test",
			"amazonec2-secret-key": "123",
			"amazonec2-region":     "us-east-1",
			"amazonec2-zone":       "e",
		},
	}

	err := driver.SetConfigFromFlags(options)

	assert.Equal(t, err, errorMissingAccessKeyOption)
}

func TestSecretKeyIsMandatory(t *testing.T) {
	driver := NewTestDriver()
	driver.awsCredentials = &cliCredentials{}
	options := &commandstest.FakeFlagger{
		Data: map[string]interface{}{
			"name":                 "test",
			"amazonec2-access-key": "foobar",
			"amazonec2-region":     "us-east-1",
			"amazonec2-zone":       "e",
		},
	}

	err := driver.SetConfigFromFlags(options)

	assert.Equal(t, err, errorMissingSecretKeyOption)
}

func TestLoadingFromCredentialsWorked(t *testing.T) {
	driver := NewCustomTestDriver(&fakeEC2WithLogin{})
	driver.awsCredentials = &fileCredentials{}
	options := &commandstest.FakeFlagger{
		Data: map[string]interface{}{
			"name":             "test",
			"amazonec2-region": "us-east-1",
			"amazonec2-zone":   "e",
		},
	}

	err := driver.SetConfigFromFlags(options)

	assert.NoError(t, err)
	assert.Equal(t, "access", driver.AccessKey)
	assert.Equal(t, "secret", driver.SecretKey)
	assert.Equal(t, "token", driver.SessionToken)
}

func TestPassingBothCLIArgWorked(t *testing.T) {
	driver := NewCustomTestDriver(&fakeEC2WithLogin{})
	driver.awsCredentials = &cliCredentials{}
	options := &commandstest.FakeFlagger{
		Data: map[string]interface{}{
			"name":                 "test",
			"amazonec2-access-key": "foobar",
			"amazonec2-secret-key": "123",
			"amazonec2-region":     "us-east-1",
			"amazonec2-zone":       "e",
		},
	}

	err := driver.SetConfigFromFlags(options)

	assert.NoError(t, err)
	assert.Equal(t, "foobar", driver.AccessKey)
	assert.Equal(t, "123", driver.SecretKey)
}
<commit_msg>Add missing test on SecurityGroupPermission<commit_after>package amazonec2

import (
	"testing"

	"errors"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/docker/machine/commands/commandstest"
	"github.com/stretchr/testify/assert"
)

const (
	testSSHPort    = 22
	testDockerPort = 2376
	testSwarmPort  = 3376
)

var (
	securityGroup = &ec2.SecurityGroup{
		GroupName: aws.String("test-group"),
		GroupId:   aws.String("12345"),
		VpcId:     aws.String("12345"),
	}
)

func TestConfigureSecurityGroupPermissionsEmpty(t *testing.T) {
	driver := NewTestDriver()

	perms := driver.configureSecurityGroupPermissions(securityGroup)

	assert.Len(t, perms, 2)
}

func TestConfigureSecurityGroupPermissionsSshOnly(t *testing.T) {
	driver := NewTestDriver()
	group := securityGroup
	group.IpPermissions = []*ec2.IpPermission{
		{
			IpProtocol: aws.String("tcp"),
			FromPort:   aws.Int64(int64(testSSHPort)),
			ToPort:     aws.Int64(int64(testSSHPort)),
		},
	}

	perms := driver.configureSecurityGroupPermissions(group)

	assert.Len(t, perms, 1)
	assert.Equal(t, testDockerPort, *perms[0].FromPort)
}

func TestConfigureSecurityGroupPermissionsDockerOnly(t *testing.T) {
	driver := NewTestDriver()
	group := securityGroup
	group.IpPermissions = []*ec2.IpPermission{
		{
			IpProtocol: aws.String("tcp"),
			FromPort:   aws.Int64((testDockerPort)),
			ToPort:     aws.Int64((testDockerPort)),
		},
	}

	perms := driver.configureSecurityGroupPermissions(group)

	assert.Len(t, perms, 1)
	assert.Equal(t, testSSHPort, *perms[0].FromPort)
}

func TestConfigureSecurityGroupPermissionsDockerAndSsh(t *testing.T) {
	driver := NewTestDriver()
	group := securityGroup
	group.IpPermissions = []*ec2.IpPermission{
		{
			IpProtocol: aws.String("tcp"),
			FromPort:   aws.Int64(testSSHPort),
			ToPort:     aws.Int64(testSSHPort),
		},
		{
			IpProtocol: aws.String("tcp"),
			FromPort:   aws.Int64(testDockerPort),
			ToPort:     aws.Int64(testDockerPort),
		},
	}

	perms := driver.configureSecurityGroupPermissions(group)

	assert.Empty(t, perms)
}

func TestConfigureSecurityGroupPermissionsWithSwarm(t *testing.T) {
	driver := NewTestDriver()
	driver.SwarmMaster = true
	group := securityGroup
	group.IpPermissions = []*ec2.IpPermission{
		{
			IpProtocol: aws.String("tcp"),
			FromPort:   aws.Int64(testSSHPort),
			ToPort:     aws.Int64(testSSHPort),
		},
		{
			IpProtocol: aws.String("tcp"),
			FromPort:   aws.Int64(testDockerPort),
			ToPort:     aws.Int64(testDockerPort),
		},
	}

	perms := driver.configureSecurityGroupPermissions(group)

	assert.Len(t, perms, 1)
	assert.Equal(t, testSwarmPort, *perms[0].FromPort)
}

func TestValidateAwsRegionValid(t *testing.T) {
	regions := []string{"eu-west-1", "eu-central-1"}

	for _, region := range regions {
		validatedRegion, err := validateAwsRegion(region)

		assert.NoError(t, err)
		assert.Equal(t, region, validatedRegion)
	}
}

func TestValidateAwsRegionInvalid(t *testing.T) {
	regions := []string{"eu-west-2", "eu-central-2"}

	for _, region := range regions {
		_, err := validateAwsRegion(region)

		assert.EqualError(t, err, "Invalid region specified")
	}
}

func TestFindDefaultVPC(t *testing.T) {
	driver := NewDriver("machineFoo", "path")
	driver.clientFactory = func() Ec2Client { return &fakeEC2WithLogin{} }

	vpc, err := driver.getDefaultVPCId()

	assert.Equal(t, "vpc-9999", vpc)
	assert.NoError(t, err)
}

func TestDefaultVPCIsMissing(t *testing.T) {
	driver := NewDriver("machineFoo", "path")
	driver.clientFactory = func() Ec2Client {
		return &fakeEC2WithDescribe{
			output: &ec2.DescribeAccountAttributesOutput{
				AccountAttributes: []*ec2.AccountAttribute{},
			},
		}
	}

	vpc, err := driver.getDefaultVPCId()

	assert.EqualError(t, err, "No default-vpc attribute")
	assert.Empty(t, vpc)
}

func TestDescribeAccountAttributeFails(t *testing.T) {
	driver := NewDriver("machineFoo", "path")
	driver.clientFactory = func() Ec2Client {
		return &fakeEC2WithDescribe{
			err: errors.New("Not Found"),
		}
	}

	vpc, err := driver.getDefaultVPCId()

	assert.EqualError(t, err, "Not Found")
	assert.Empty(t, vpc)
}

func TestAccessKeyIsMandatory(t *testing.T) {
	driver := NewTestDriver()
	driver.awsCredentials = &cliCredentials{}
	options := &commandstest.FakeFlagger{
		Data: map[string]interface{}{
			"name":             "test",
			"amazonec2-region": "us-east-1",
			"amazonec2-zone":   "e",
		},
	}

	err := driver.SetConfigFromFlags(options)

	assert.Equal(t, err, errorMissingAccessKeyOption)
}

func TestAccessKeyIsMandatoryEvenIfSecretKeyIsPassed(t *testing.T) {
	driver := NewTestDriver()
	driver.awsCredentials = &cliCredentials{}
	options := &commandstest.FakeFlagger{
		Data: map[string]interface{}{
			"name":                 "test",
			"amazonec2-secret-key": "123",
			"amazonec2-region":     "us-east-1",
			"amazonec2-zone":       "e",
		},
	}

	err := driver.SetConfigFromFlags(options)

	assert.Equal(t, err, errorMissingAccessKeyOption)
}

func TestSecretKeyIsMandatory(t *testing.T) {
	driver := NewTestDriver()
	driver.awsCredentials = &cliCredentials{}
	options := &commandstest.FakeFlagger{
		Data: map[string]interface{}{
			"name":                 "test",
			"amazonec2-access-key": "foobar",
			"amazonec2-region":     "us-east-1",
			"amazonec2-zone":       "e",
		},
	}

	err := driver.SetConfigFromFlags(options)

	assert.Equal(t, err, errorMissingSecretKeyOption)
}

func TestLoadingFromCredentialsWorked(t *testing.T) {
	driver := NewCustomTestDriver(&fakeEC2WithLogin{})
	driver.awsCredentials = &fileCredentials{}
	options := &commandstest.FakeFlagger{
		Data: map[string]interface{}{
			"name":             "test",
			"amazonec2-region": "us-east-1",
			"amazonec2-zone":   "e",
		},
	}

	err := driver.SetConfigFromFlags(options)

	assert.NoError(t, err)
	assert.Equal(t, "access", driver.AccessKey)
	assert.Equal(t, "secret", driver.SecretKey)
	assert.Equal(t, "token", driver.SessionToken)
}

func TestPassingBothCLIArgWorked(t *testing.T) {
	driver := NewCustomTestDriver(&fakeEC2WithLogin{})
	driver.awsCredentials = &cliCredentials{}
	options := &commandstest.FakeFlagger{
		Data: map[string]interface{}{
			"name":                 "test",
			"amazonec2-access-key": "foobar",
			"amazonec2-secret-key": "123",
			"amazonec2-region":     "us-east-1",
			"amazonec2-zone":       "e",
		},
	}

	err := driver.SetConfigFromFlags(options)

	assert.NoError(t, err)
	assert.Equal(t, "foobar", driver.AccessKey)
	assert.Equal(t, "123", driver.SecretKey)
}
<|endoftext|>


<commit_before>package sequence

/*
Note :
(1) store the sequence in the ETCD cluster, and local file(sequence.dat)
(2) batch get the sequences from ETCD cluster, and store the max sequence id in the local file
(3) the sequence range is : [currentSeqId, maxSeqId), when the currentSeqId >= maxSeqId, fetch the new maxSeqId.
*/

import (
	"context"
	"fmt"
	"sync"
	"time"

	"io"
	"os"
	"strconv"
	"strings"

	"github.com/chrislusf/seaweedfs/weed/glog"
	"go.etcd.io/etcd/client"
)

const (
	// EtcdKeyPrefix                   = "/seaweedfs"
	EtcdKeySequence                 = "/master/sequence"
	EtcdContextTimeoutSecond        = 100 * time.Second
	DefaultEtcdSteps         uint64 = 500 // internal counter
	SequencerFileName               = "sequencer.dat"
	FileMaxSequenceLength           = 128
)

type EtcdSequencer struct {
	sequenceLock sync.Mutex

	// available sequence range : [currentSeqId, maxSeqId)
	currentSeqId uint64
	maxSeqId     uint64

	keysAPI client.KeysAPI
	seqFile *os.File
}

func NewEtcdSequencer(etcdUrls string, metaFolder string) (*EtcdSequencer, error) {
	file, err := openSequenceFile(metaFolder + "/" + SequencerFileName)
	if nil != err {
		return nil, fmt.Errorf("open sequence file fialed, %v", err)
	}
	defer file.Close()

	cli, err := client.New(client.Config{
		Endpoints: strings.Split(etcdUrls, ","),
		Username:  "",
		Password:  "",
	})
	if err != nil {
		return nil, err
	}
	keysApi := client.NewKeysAPI(cli)

	// TODO: the current sequence id in local file is not used
	maxValue, _, err := readSequenceFile(file)
	if err != nil {
		return nil, fmt.Errorf("read sequence from file failed, %v", err)
	}
	glog.V(4).Infof("read sequence from file : %d", maxValue)

	newSeq, err := setMaxSequenceToEtcd(keysApi, maxValue)
	if err != nil {
		return nil, err
	}

	sequencer := &EtcdSequencer{maxSeqId: newSeq,
		currentSeqId: newSeq,
		keysAPI:      keysApi,
		seqFile:      file,
	}
	return sequencer, nil
}

func (es *EtcdSequencer) NextFileId(count uint64) uint64 {
	es.sequenceLock.Lock()
	defer es.sequenceLock.Unlock()

	if (es.currentSeqId + count) >= es.maxSeqId {
		reqSteps := DefaultEtcdSteps
		if count > DefaultEtcdSteps {
			reqSteps += count
		}
		maxId, err := batchGetSequenceFromEtcd(es.keysAPI, reqSteps)
		glog.V(4).Infof("get max sequence id from etcd, %d", maxId)
		if err != nil {
			glog.Error(err)
			return 0
		}
		es.currentSeqId, es.maxSeqId = maxId-reqSteps, maxId
		glog.V(4).Infof("current id : %d, max id : %d", es.currentSeqId, es.maxSeqId)

		if err := writeSequenceFile(es.seqFile, es.maxSeqId, es.currentSeqId); err != nil {
			glog.Errorf("flush sequence to file failed, %v", err)
		}
	}

	ret := es.currentSeqId
	es.currentSeqId += count
	return ret
}

/**
instead of collecting the max value from volume server,
the max value should be saved in local config file and ETCD cluster
*/
func (es *EtcdSequencer) SetMax(seenValue uint64) {
	es.sequenceLock.Lock()
	defer es.sequenceLock.Unlock()
	if seenValue > es.maxSeqId {
		maxId, err := setMaxSequenceToEtcd(es.keysAPI, seenValue)
		if err != nil {
			glog.Errorf("set Etcd Max sequence failed : %v", err)
			return
		}
		es.currentSeqId, es.maxSeqId = maxId, maxId

		if err := writeSequenceFile(es.seqFile, maxId, maxId); err != nil {
			glog.Errorf("flush sequence to file failed, %v", err)
		}
	}
}

func (es *EtcdSequencer) GetMax() uint64 {
	return es.maxSeqId
}

func (es *EtcdSequencer) Peek() uint64 {
	return es.currentSeqId
}

func batchGetSequenceFromEtcd(kvApi client.KeysAPI, step uint64) (uint64, error) {
	if step <= 0 {
		return 0, fmt.Errorf("the step must be large than 1")
	}

	ctx, cancel := context.WithTimeout(context.Background(), EtcdContextTimeoutSecond)
	var endSeqValue uint64 = 0
	defer cancel()
	for {
		getResp, err := kvApi.Get(ctx, EtcdKeySequence, &client.GetOptions{Recursive: false, Quorum: true})
		if err != nil {
			return 0, err
		}
		if getResp.Node == nil {
			continue
		}

		prevValue := getResp.Node.Value
		prevSeqValue, err := strconv.ParseUint(prevValue, 10, 64)
		if err != nil {
			return 0, fmt.Errorf("get sequence from etcd failed, %v", err)
		}
		endSeqValue = prevSeqValue + step
		endSeqStr := strconv.FormatUint(endSeqValue, 10)

		_, err = kvApi.Set(ctx, EtcdKeySequence, endSeqStr, &client.SetOptions{PrevValue: prevValue})
		if err == nil {
			break
		}
		glog.Error(err)
	}

	return endSeqValue, nil
}

/**
update the value of the key EtcdKeySequence in ETCD cluster with the parameter of maxSeq,
when the value of the key EtcdKeySequence is equal to or large than the parameter maxSeq,
return the value of EtcdKeySequence in the ETCD cluster;
when the value of the EtcdKeySequence is less than the parameter maxSeq,
return the value of the parameter maxSeq
*/
func setMaxSequenceToEtcd(kvApi client.KeysAPI, maxSeq uint64) (uint64, error) {
	maxSeqStr := strconv.FormatUint(maxSeq, 10)
	ctx, cancel := context.WithTimeout(context.Background(), EtcdContextTimeoutSecond)
	defer cancel()

	for {
		getResp, err := kvApi.Get(ctx, EtcdKeySequence, &client.GetOptions{Recursive: false, Quorum: true})
		if err != nil {
			if ce, ok := err.(client.Error); ok && (ce.Code == client.ErrorCodeKeyNotFound) {
				_, err := kvApi.Create(ctx, EtcdKeySequence, maxSeqStr)
				if err == nil {
					continue
				}
				if ce, ok = err.(client.Error); ok && (ce.Code == client.ErrorCodeNodeExist) {
					continue
				}
				return 0, err
			} else {
				return 0, err
			}
		}

		if getResp.Node == nil {
			continue
		}
		prevSeqStr := getResp.Node.Value
		prevSeq, err := strconv.ParseUint(prevSeqStr, 10, 64)
		if err != nil {
			return 0, err
		}
		if prevSeq >= maxSeq {
			return prevSeq, nil
		}

		_, err = kvApi.Set(ctx, EtcdKeySequence, maxSeqStr, &client.SetOptions{PrevValue: prevSeqStr})
		if err != nil {
			return 0, err
		}
	}
}

func openSequenceFile(file string) (*os.File, error) {
	_, err := os.Stat(file)
	if os.IsNotExist(err) {
		fid, err := os.OpenFile(file, os.O_RDWR|os.O_CREATE, 0644)
		if err != nil {
			return nil, err
		}
		if err := writeSequenceFile(fid, 1, 0); err != nil {
			return nil, err
		}
		return fid, nil
	} else {
		return os.OpenFile(file, os.O_RDWR|os.O_CREATE, 0644)
	}
}

/*
read sequence and step from sequence file
*/
func readSequenceFile(file *os.File) (uint64, uint64, error) {
	sequence := make([]byte, FileMaxSequenceLength)
	size, err := file.ReadAt(sequence, 0)
	if (err != nil) && (err != io.EOF) {
		err := fmt.Errorf("cannot read file %s, %v", file.Name(), err)
		return 0, 0, err
	}
	sequence = sequence[0:size]
	seqs := strings.Split(string(sequence), ":")
	maxId, err := strconv.ParseUint(seqs[0], 10, 64)
	if err != nil {
		return 0, 0, fmt.Errorf("parse sequence from file failed, %v", err)
	}

	if len(seqs) > 1 {
		step, err := strconv.ParseUint(seqs[1], 10, 64)
		if err != nil {
			return 0, 0, fmt.Errorf("parse sequence from file failed, %v", err)
		}
		return maxId, step, nil
	}

	return maxId, 0, nil
}

/**
write the sequence and step to sequence file
*/
func writeSequenceFile(file *os.File, sequence, step uint64) error {
	_ = step
	seqStr := fmt.Sprintf("%d:%d", sequence, sequence)
	if _, err := file.Seek(0, 0); err != nil {
		err = fmt.Errorf("cannot seek to the beginning of %s: %v", file.Name(), err)
		return err
	}
	if err := file.Truncate(0); err != nil {
		return fmt.Errorf("truncate sequence file faield : %v", err)
	}
	if _, err := file.WriteString(seqStr); err != nil {
		return fmt.Errorf("write file %s failed, %v", file.Name(), err)
	}
	if err := file.Sync(); err != nil {
		return fmt.Errorf("flush file %s failed, %v", file.Name(), err)
	}
	return nil
}

// the UT helper method
// func deleteEtcdKey(kvApi client.KeysAPI, key string) error {
// 	ctx, cancel := context.WithTimeout(context.Background(), EtcdContextTimeoutSecond)
// 	defer cancel()
// 	_, err := kvApi.Delete(ctx, key, &client.DeleteOptions{Dir: false})
// 	if err != nil {
// 		return err
// 	}
// 	return nil
// }
<commit_msg>Revert "close opened file"<commit_after>package sequence

/*
Note :
(1) store the sequence in the ETCD cluster, and local file(sequence.dat)
(2) batch get the sequences from ETCD cluster, and store the max sequence id in the local file
(3) the sequence range is : [currentSeqId, maxSeqId), when the currentSeqId >= maxSeqId, fetch the new maxSeqId.
*/

import (
	"context"
	"fmt"
	"sync"
	"time"

	"io"
	"os"
	"strconv"
	"strings"

	"github.com/chrislusf/seaweedfs/weed/glog"
	"go.etcd.io/etcd/client"
)

const (
	// EtcdKeyPrefix                   = "/seaweedfs"
	EtcdKeySequence                 = "/master/sequence"
	EtcdContextTimeoutSecond        = 100 * time.Second
	DefaultEtcdSteps         uint64 = 500 // internal counter
	SequencerFileName               = "sequencer.dat"
	FileMaxSequenceLength           = 128
)

type EtcdSequencer struct {
	sequenceLock sync.Mutex

	// available sequence range : [currentSeqId, maxSeqId)
	currentSeqId uint64
	maxSeqId     uint64

	keysAPI client.KeysAPI
	seqFile *os.File
}

func NewEtcdSequencer(etcdUrls string, metaFolder string) (*EtcdSequencer, error) {
	file, err := openSequenceFile(metaFolder + "/" + SequencerFileName)
	if nil != err {
		return nil, fmt.Errorf("open sequence file fialed, %v", err)
	}

	cli, err := client.New(client.Config{
		Endpoints: strings.Split(etcdUrls, ","),
		Username:  "",
		Password:  "",
	})
	if err != nil {
		return nil, err
	}
	keysApi := client.NewKeysAPI(cli)

	// TODO: the current sequence id in local file is not used
	maxValue, _, err := readSequenceFile(file)
	if err != nil {
		return nil, fmt.Errorf("read sequence from file failed, %v", err)
	}
	glog.V(4).Infof("read sequence from file : %d", maxValue)

	newSeq, err := setMaxSequenceToEtcd(keysApi, maxValue)
	if err != nil {
		return nil, err
	}

	sequencer := &EtcdSequencer{maxSeqId: newSeq,
		currentSeqId: newSeq,
		keysAPI:      keysApi,
		seqFile:      file,
	}
	return sequencer, nil
}

func (es *EtcdSequencer) NextFileId(count uint64) uint64 {
	es.sequenceLock.Lock()
	defer es.sequenceLock.Unlock()

	if (es.currentSeqId + count) >= es.maxSeqId {
		reqSteps := DefaultEtcdSteps
		if count > DefaultEtcdSteps {
			reqSteps += count
		}
		maxId, err := batchGetSequenceFromEtcd(es.keysAPI, reqSteps)
		glog.V(4).Infof("get max sequence id from etcd, %d", maxId)
		if err != nil {
			glog.Error(err)
			return 0
		}
		es.currentSeqId, es.maxSeqId = maxId-reqSteps, maxId
		glog.V(4).Infof("current id : %d, max id : %d", es.currentSeqId, es.maxSeqId)

		if err := writeSequenceFile(es.seqFile, es.maxSeqId, es.currentSeqId); err != nil {
			glog.Errorf("flush sequence to file failed, %v", err)
		}
	}

	ret := es.currentSeqId
	es.currentSeqId += count
	return ret
}

/**
instead of collecting the max value from volume server,
the max value should be saved in local config file and ETCD cluster
*/
func (es *EtcdSequencer) SetMax(seenValue uint64) {
	es.sequenceLock.Lock()
	defer es.sequenceLock.Unlock()
	if seenValue > es.maxSeqId {
		maxId, err := setMaxSequenceToEtcd(es.keysAPI, seenValue)
		if err != nil {
			glog.Errorf("set Etcd Max sequence failed : %v", err)
			return
		}
		es.currentSeqId, es.maxSeqId = maxId, maxId

		if err := writeSequenceFile(es.seqFile, maxId, maxId); err != nil {
			glog.Errorf("flush sequence to file failed, %v", err)
		}
	}
}

func (es *EtcdSequencer) GetMax() uint64 {
	return es.maxSeqId
}

func (es *EtcdSequencer) Peek() uint64 {
	return es.currentSeqId
}

func batchGetSequenceFromEtcd(kvApi client.KeysAPI, step uint64) (uint64, error) {
	if step <= 0 {
		return 0, fmt.Errorf("the step must be large than 1")
	}

	ctx, cancel := context.WithTimeout(context.Background(), EtcdContextTimeoutSecond)
	var endSeqValue uint64 = 0
	defer cancel()
	for {
		getResp, err := kvApi.Get(ctx, EtcdKeySequence, &client.GetOptions{Recursive: false, Quorum: true})
		if err != nil {
			return 0, err
		}
		if getResp.Node == nil {
			continue
		}

		prevValue := getResp.Node.Value
		prevSeqValue, err := strconv.ParseUint(prevValue, 10, 64)
		if err != nil {
			return 0, fmt.Errorf("get sequence from etcd failed, %v", err)
		}
		endSeqValue = prevSeqValue + step
		endSeqStr := strconv.FormatUint(endSeqValue, 10)

		_, err = kvApi.Set(ctx, EtcdKeySequence, endSeqStr, &client.SetOptions{PrevValue: prevValue})
		if err == nil {
			break
		}
		glog.Error(err)
	}

	return endSeqValue, nil
}

/**
update the value of the key EtcdKeySequence in ETCD cluster with the parameter of maxSeq,
when the value of the key EtcdKeySequence is equal to or large than the parameter maxSeq,
return the value of EtcdKeySequence in the ETCD cluster;
when the value of the EtcdKeySequence is less than the parameter maxSeq,
return the value of the parameter maxSeq
*/
func setMaxSequenceToEtcd(kvApi client.KeysAPI, maxSeq uint64) (uint64, error) {
	maxSeqStr := strconv.FormatUint(maxSeq, 10)
	ctx, cancel := context.WithTimeout(context.Background(), EtcdContextTimeoutSecond)
	defer cancel()

	for {
		getResp, err := kvApi.Get(ctx, EtcdKeySequence, &client.GetOptions{Recursive: false, Quorum: true})
		if err != nil {
			if ce, ok := err.(client.Error); ok && (ce.Code == client.ErrorCodeKeyNotFound) {
				_, err := kvApi.Create(ctx, EtcdKeySequence, maxSeqStr)
				if err == nil {
					continue
				}
				if ce, ok = err.(client.Error); ok && (ce.Code == client.ErrorCodeNodeExist) {
					continue
				}
				return 0, err
			} else {
				return 0, err
			}
		}

		if getResp.Node == nil {
			continue
		}
		prevSeqStr := getResp.Node.Value
		prevSeq, err := strconv.ParseUint(prevSeqStr, 10, 64)
		if err != nil {
			return 0, err
		}
		if prevSeq >= maxSeq {
			return prevSeq, nil
		}

		_, err = kvApi.Set(ctx, EtcdKeySequence, maxSeqStr, &client.SetOptions{PrevValue: prevSeqStr})
		if err != nil {
			return 0, err
		}
	}
}

func openSequenceFile(file string) (*os.File, error) {
	_, err := os.Stat(file)
	if os.IsNotExist(err) {
		fid, err := os.OpenFile(file, os.O_RDWR|os.O_CREATE, 0644)
		if err != nil {
			return nil, err
		}
		if err := writeSequenceFile(fid, 1, 0); err != nil {
			return nil, err
		}
		return fid, nil
	} else {
		return os.OpenFile(file, os.O_RDWR|os.O_CREATE, 0644)
	}
}

/*
read sequence and step from sequence file
*/
func readSequenceFile(file *os.File) (uint64, uint64, error) {
	sequence := make([]byte, FileMaxSequenceLength)
	size, err := file.ReadAt(sequence, 0)
	if (err != nil) && (err != io.EOF) {
		err := fmt.Errorf("cannot read file %s, %v", file.Name(), err)
		return 0, 0, err
	}
	sequence = sequence[0:size]
	seqs := strings.Split(string(sequence), ":")
	maxId, err := strconv.ParseUint(seqs[0], 10, 64)
	if err != nil {
		return 0, 0, fmt.Errorf("parse sequence from file failed, %v", err)
	}

	if len(seqs) > 1 {
		step, err := strconv.ParseUint(seqs[1], 10, 64)
		if err != nil {
			return 0, 0, fmt.Errorf("parse sequence from file failed, %v", err)
		}
		return maxId, step, nil
	}

	return maxId, 0, nil
}

/**
write the sequence and step to sequence file
*/
func writeSequenceFile(file *os.File, sequence, step uint64) error {
	_ = step
	seqStr := fmt.Sprintf("%d:%d", sequence, sequence)
	if _, err := file.Seek(0, 0); err != nil {
		err = fmt.Errorf("cannot seek to the beginning of %s: %v", file.Name(), err)
		return err
	}
	if err := file.Truncate(0); err != nil {
		return fmt.Errorf("truncate sequence file faield : %v", err)
	}
	if _, err := file.WriteString(seqStr); err != nil {
		return fmt.Errorf("write file %s failed, %v", file.Name(), err)
	}
	if err := file.Sync(); err != nil {
		return fmt.Errorf("flush file %s failed, %v", file.Name(), err)
	}
	return nil
}

// the UT helper method
// func deleteEtcdKey(kvApi client.KeysAPI, key string) error {
// 	ctx, cancel := context.WithTimeout(context.Background(), EtcdContextTimeoutSecond)
// 	defer cancel()
// 	_, err := kvApi.Delete(ctx, key, &client.DeleteOptions{Dir: false})
// 	if err != nil {
// 		return err
// 	}
// 	return nil
// }
<|endoftext|>


<commit_before>package cache

import (
	"context"
	"flag"
	"runtime"
	"sync"

	"github.com/grafana/metrictank/mdata/cache/accnt"
	"github.com/grafana/metrictank/mdata/chunk"
	"github.com/grafana/metrictank/stats"
	"github.com/grafana/metrictank/tracing"
	opentracing "github.com/opentracing/opentracing-go"
	"github.com/raintank/worldping-api/pkg/log"
	"github.com/rakyll/globalconf"
)

var (
	maxSize      uint64
	searchFwdBug = stats.NewCounter32("recovered_errors.cache.metric.searchForwardBug")
)

func init() {
	flags := flag.NewFlagSet("chunk-cache", flag.ExitOnError)
	// (1024 ^ 3) * 4 = 4294967296 = 4G
	flags.Uint64Var(&maxSize, "max-size", 4294967296, "Maximum size of chunk cache in bytes")
	globalconf.Register("chunk-cache", flags)
}

type CCache struct {
	sync.RWMutex

	// one CCacheMetric struct per metric key, indexed by the key
	metricCache map[string]*CCacheMetric

	// sets of metric keys, indexed by their raw metric keys
	metricRawKeys map[string]map[string]struct{}

	// accounting for the cache. keeps track of when data needs to be evicted
	// and what should be evicted
	accnt accnt.Accnt

	// channel that's only used to signal go routines to stop
	stop chan interface{}

	tracer opentracing.Tracer
}

func NewCCache() *CCache {
	cc := &CCache{
		metricCache:   make(map[string]*CCacheMetric),
		metricRawKeys: make(map[string]map[string]struct{}),
		accnt:         accnt.NewFlatAccnt(maxSize),
		stop:          make(chan interface{}),
		tracer:        opentracing.NoopTracer{},
	}
	go cc.evictLoop()
	return cc
}

func (c *CCache) SetTracer(t opentracing.Tracer) {
	c.tracer = t
}

func (c *CCache) evictLoop() {
	evictQ := c.accnt.GetEvictQ()
	for {
		select {
		case target := <-evictQ:
			c.evict(target)
		case _ = <-c.stop:
			return
		}
	}
}

// takes a raw key and deletes all archives associated with it from cache
func (c *CCache) DelMetric(rawMetric string) (int, int) {
	archives, series := 0, 0

	c.Lock()
	defer c.Unlock()

	mets, ok := c.metricRawKeys[rawMetric]
	if !ok {
		return archives, series
	}

	for met := range mets {
		delete(c.metricCache, met)
		c.accnt.DelMetric(met)
		archives++
	}

	delete(c.metricRawKeys, rawMetric)
	series++

	return series, archives
}

// adds the given chunk to the cache, but only if the metric is sufficiently hot
func (c *CCache) CacheIfHot(metric string, prev uint32, itergen chunk.IterGen) {
	c.RLock()

	var met *CCacheMetric
	var ok bool

	// if this metric is not cached at all it is not hot
	if met, ok = c.metricCache[metric]; !ok {
		c.RUnlock()
		return
	}

	// if the previous chunk is not cached we consider the metric not hot enough to cache this chunk
	// only works reliably if the last chunk of that metric is span aware, otherwise lastTs() will be guessed
	// conservatively which means that the returned value will probably be lower than the real last ts
	if met.lastTs() < itergen.Ts {
		c.RUnlock()
		return
	}

	accnt.CacheChunkPushHot.Inc()

	c.RUnlock()
	met.Add(prev, itergen)
}

func (c *CCache) Add(metric, rawMetric string, prev uint32, itergen chunk.IterGen) {
	c.Lock()
	defer c.Unlock()

	ccm, ok := c.metricCache[metric]
	if !ok {
		ccm = NewCCacheMetric()
		ccm.Init(rawMetric, prev, itergen)
		c.metricCache[metric] = ccm

		// if we do not have this raw key yet, create an entry for it
		ccms, ok := c.metricRawKeys[rawMetric]
		if !ok {
			ccms = make(map[string]struct{})
			c.metricRawKeys[rawMetric] = ccms
		}

		// if we don't have it yet, associate the metric with this raw key
		if _, ok = ccms[metric]; !ok {
			ccms[metric] = struct{}{}
		}
	} else {
		ccm.Add(prev, itergen)
	}

	c.accnt.AddChunk(metric, itergen.Ts, itergen.Size())
}

func (cc *CCache) Reset() (int, int) {
	cc.Lock()
	cc.accnt.Reset()
	series := len(cc.metricRawKeys)
	archives := len(cc.metricCache)
	cc.metricCache = make(map[string]*CCacheMetric)
	cc.metricRawKeys = make(map[string]map[string]struct{})
	cc.Unlock()
	return series, archives
}

func (c *CCache) Stop() {
	c.accnt.Stop()
	c.stop <- nil
}

func (c *CCache) evict(target *accnt.EvictTarget) {
	c.Lock()
	// evict() might get called many times in a loop, but we don't want it to block
	// cache reads with the write lock, so we yield right after unlocking to allow
	// reads to go first.
	defer runtime.Gosched()
	defer c.Unlock()

	ccm, ok := c.metricCache[target.Metric]
	if !ok {
		return
	}

	log.Debug("CCache evict: evicting chunk %d on metric %s\n", target.Ts, target.Metric)
	length := c.metricCache[target.Metric].Del(target.Ts)
	if length == 0 {
		delete(c.metricCache, target.Metric)

		// this key should alway be present, if not there there is a corruption of the state
		delete(c.metricRawKeys[ccm.RawMetric], target.Metric)
		if len(c.metricRawKeys[ccm.RawMetric]) == 0 {
			delete(c.metricRawKeys, ccm.RawMetric)
		}
	}
}

func (c *CCache) Search(ctx context.Context, metric string, from, until uint32) *CCSearchResult {
	ctx, span := tracing.NewSpan(ctx, c.tracer, "CCache.Search")
	defer span.Finish()
	var hit chunk.IterGen
	var cm *CCacheMetric
	var ok bool
	res := &CCSearchResult{
		From:  from,
		Until: until,
	}

	if from == until {
		return res
	}

	c.RLock()
	defer c.RUnlock()

	if cm, ok = c.metricCache[metric]; !ok {
		span.SetTag("cache", "miss")
		accnt.CacheMetricMiss.Inc()
		return res
	}

	cm.Search(ctx, metric, res, from, until)
	if len(res.Start) == 0 && len(res.End) == 0 {
		span.SetTag("cache", "miss")
		accnt.CacheMetricMiss.Inc()
	} else {

		accnt.CacheChunkHit.Add(len(res.Start) + len(res.End))
		go func() {
			for _, hit = range res.Start {
				c.accnt.HitChunk(metric, hit.Ts)
			}
			for _, hit = range res.End {
				c.accnt.HitChunk(metric, hit.Ts)
			}
		}()

		if res.Complete {
			span.SetTag("cache", "hit-full")
			accnt.CacheMetricHitFull.Inc()
		} else {
			span.SetTag("cache", "hit-partial")
			accnt.CacheMetricHitPartial.Inc()
		}
	}

	return res
}
<commit_msg>small CCache.Add optimization<commit_after>package cache

import (
	"context"
	"flag"
	"runtime"
	"sync"

	"github.com/grafana/metrictank/mdata/cache/accnt"
	"github.com/grafana/metrictank/mdata/chunk"
	"github.com/grafana/metrictank/stats"
	"github.com/grafana/metrictank/tracing"
	opentracing "github.com/opentracing/opentracing-go"
	"github.com/raintank/worldping-api/pkg/log"
	"github.com/rakyll/globalconf"
)

var (
	maxSize      uint64
	searchFwdBug = stats.NewCounter32("recovered_errors.cache.metric.searchForwardBug")
)

func init() {
	flags := flag.NewFlagSet("chunk-cache", flag.ExitOnError)
	// (1024 ^ 3) * 4 = 4294967296 = 4G
	flags.Uint64Var(&maxSize, "max-size", 4294967296, "Maximum size of chunk cache in bytes")
	globalconf.Register("chunk-cache", flags)
}

type CCache struct {
	sync.RWMutex

	// one CCacheMetric struct per metric key, indexed by the key
	metricCache map[string]*CCacheMetric

	// sets of metric keys, indexed by their raw metric keys
	metricRawKeys map[string]map[string]struct{}

	// accounting for the cache. keeps track of when data needs to be evicted
	// and what should be evicted
	accnt accnt.Accnt

	// channel that's only used to signal go routines to stop
	stop chan interface{}

	tracer opentracing.Tracer
}

func NewCCache() *CCache {
	cc := &CCache{
		metricCache:   make(map[string]*CCacheMetric),
		metricRawKeys: make(map[string]map[string]struct{}),
		accnt:         accnt.NewFlatAccnt(maxSize),
		stop:          make(chan interface{}),
		tracer:        opentracing.NoopTracer{},
	}
	go cc.evictLoop()
	return cc
}

func (c *CCache) SetTracer(t opentracing.Tracer) {
	c.tracer = t
}

func (c *CCache) evictLoop() {
	evictQ := c.accnt.GetEvictQ()
	for {
		select {
		case target := <-evictQ:
			c.evict(target)
		case _ = <-c.stop:
			return
		}
	}
}

// takes a raw key and deletes all archives associated with it from cache
func (c *CCache) DelMetric(rawMetric string) (int, int) {
	archives, series := 0, 0

	c.Lock()
	defer c.Unlock()

	mets, ok := c.metricRawKeys[rawMetric]
	if !ok {
		return archives, series
	}

	for met := range mets {
		delete(c.metricCache, met)
		c.accnt.DelMetric(met)
		archives++
	}

	delete(c.metricRawKeys, rawMetric)
	series++

	return series, archives
}

// adds the given chunk to the cache, but only if the metric is sufficiently hot
func (c *CCache) CacheIfHot(metric string, prev uint32, itergen chunk.IterGen) {
	c.RLock()

	var met *CCacheMetric
	var ok bool

	// if this metric is not cached at all it is not hot
	if met, ok = c.metricCache[metric]; !ok {
		c.RUnlock()
		return
	}

	// if the previous chunk is not cached we consider the metric not hot enough to cache this chunk
	// only works reliably if the last chunk of that metric is span aware, otherwise lastTs() will be guessed
	// conservatively which means that the returned value will probably be lower than the real last ts
	if met.lastTs() < itergen.Ts {
		c.RUnlock()
		return
	}

	accnt.CacheChunkPushHot.Inc()

	c.RUnlock()
	met.Add(prev, itergen)
}

func (c *CCache) Add(metric, rawMetric string, prev uint32, itergen chunk.IterGen) {
	c.Lock()
	defer c.Unlock()

	ccm, ok := c.metricCache[metric]
	if !ok {
		ccm = NewCCacheMetric()
		ccm.Init(rawMetric, prev, itergen)
		c.metricCache[metric] = ccm

		// if we do not have this raw key yet, create the entry with the association
		ccms, ok := c.metricRawKeys[rawMetric]
		if !ok {
			c.metricRawKeys[rawMetric] = map[string]struct{}{
				metric: {},
			}
		} else {
			// otherwise, make sure the association exists
			ccms[metric] = struct{}{}
		}
	} else {
		ccm.Add(prev, itergen)
	}

	c.accnt.AddChunk(metric, itergen.Ts, itergen.Size())
}

func (cc *CCache) Reset() (int, int) {
	cc.Lock()
	cc.accnt.Reset()
	series := len(cc.metricRawKeys)
	archives := len(cc.metricCache)
	cc.metricCache = make(map[string]*CCacheMetric)
	cc.metricRawKeys = make(map[string]map[string]struct{})
	cc.Unlock()
	return series, archives
}

func (c *CCache) Stop() {
	c.accnt.Stop()
	c.stop <- nil
}

func (c *CCache) evict(target *accnt.EvictTarget) {
	c.Lock()
	// evict() might get called many times in a loop, but we don't want it to block
	// cache reads with the write lock, so we yield right after unlocking to allow
	// reads to go first.
	defer runtime.Gosched()
	defer c.Unlock()

	ccm, ok := c.metricCache[target.Metric]
	if !ok {
		return
	}

	log.Debug("CCache evict: evicting chunk %d on metric %s\n", target.Ts, target.Metric)
	length := c.metricCache[target.Metric].Del(target.Ts)
	if length == 0 {
		delete(c.metricCache, target.Metric)

		// this key should alway be present, if not there there is a corruption of the state
		delete(c.metricRawKeys[ccm.RawMetric], target.Metric)
		if len(c.metricRawKeys[ccm.RawMetric]) == 0 {
			delete(c.metricRawKeys, ccm.RawMetric)
		}
	}
}

func (c *CCache) Search(ctx context.Context, metric string, from, until uint32) *CCSearchResult {
	ctx, span := tracing.NewSpan(ctx, c.tracer, "CCache.Search")
	defer span.Finish()
	var hit chunk.IterGen
	var cm *CCacheMetric
	var ok bool
	res := &CCSearchResult{
		From:  from,
		Until: until,
	}

	if from == until {
		return res
	}

	c.RLock()
	defer c.RUnlock()

	if cm, ok = c.metricCache[metric]; !ok {
		span.SetTag("cache", "miss")
		accnt.CacheMetricMiss.Inc()
		return res
	}

	cm.Search(ctx, metric, res, from, until)
	if len(res.Start) == 0 && len(res.End) == 0 {
		span.SetTag("cache", "miss")
		accnt.CacheMetricMiss.Inc()
	} else {

		accnt.CacheChunkHit.Add(len(res.Start) + len(res.End))
		go func() {
			for _, hit = range res.Start {
				c.accnt.HitChunk(metric, hit.Ts)
			}
			for _, hit = range res.End {
				c.accnt.HitChunk(metric, hit.Ts)
			}
		}()

		if res.Complete {
			span.SetTag("cache", "hit-full")
			accnt.CacheMetricHitFull.Inc()
		} else {
			span.SetTag("cache", "hit-partial")
			accnt.CacheMetricHitPartial.Inc()
		}
	}

	return res
}
<|endoftext|>


<commit_before>// Package types is where exportable API structures go. This is so we can share
// the XML marshalling and unmarshalling with the services.
package types

import (
	"encoding/xml"
	"time"
)

type Version struct {
	XMLName xml.Name `xml:"version"`
	App     *App     `xml:"app"`
	Package *Package `xml:"package"`
}

type App struct {
	XMLName  xml.Name  `xml:"app"`
	Id       string    `xml:"id,attr"`
	Version  string    `xml:"version,attr"`
	Track    string    `xml:"track,attr"`
	Date     time.Time `xml:"-"`
	IsActive bool      `xml:"-"`
}

type Package struct {
	XMLName              xml.Name `xml:"package"`
	Name                 string   `xml:"name,attr"`      // Package filename
	Size                 string   `xml:"size,attr"`      // Size of the file (in bytes)
	Path                 string   `xml:"path,attr"`      // Path from the root to the file
	Sha1Sum              string   `xml:"sha1sum,attr"`   // SHA-1 hash of the file
	Sha256Sum            string   `xml:"sha256sum,attr"` // Sha-256 hash of the file (extension)
	Required             bool     `xml:"required,attr"`
	MetadataSignatureRsa string   `xml:"MetadataSignatureRsa,attr,omitempty" json:"metadata_signature_rsa"`
	MetadataSize         string   `xml:"MetadataSize,attr,omitempty" json:"metadata_size"`
}
<commit_msg>feat(update/types): nil datastore tag on XMLName<commit_after>// Package types is where exportable API structures go. This is so we can share
// the XML marshalling and unmarshalling with the services.
package types

import (
	"encoding/xml"
	"time"
)

type Version struct {
	XMLName xml.Name `xml:"version" datastore:"-"`
	App     *App     `xml:"app"`
	Package *Package `xml:"package"`
}

type App struct {
	XMLName  xml.Name  `xml:"app" datastore:"-"`
	Id       string    `xml:"id,attr"`
	Version  string    `xml:"version,attr"`
	Track    string    `xml:"track,attr"`
	Date     time.Time `xml:"-"`
	IsActive bool      `xml:"-"`
}

type Package struct {
	XMLName              xml.Name `xml:"package" datastore:"-"`
	Name                 string   `xml:"name,attr"`      // Package filename
	Size                 string   `xml:"size,attr"`      // Size of the file (in bytes)
	Path                 string   `xml:"path,attr"`      // Path from the root to the file
	Sha1Sum              string   `xml:"sha1sum,attr"`   // SHA-1 hash of the file
	Sha256Sum            string   `xml:"sha256sum,attr"` // Sha-256 hash of the file (extension)
	Required             bool     `xml:"required,attr"`
	MetadataSignatureRsa string   `xml:"MetadataSignatureRsa,attr,omitempty" json:"metadata_signature_rsa"`
	MetadataSize         string   `xml:"MetadataSize,attr,omitempty" json:"metadata_size"`
}
<|endoftext|>


<commit_before>package util

import (
	"testing"

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

func TestCompressString(t *testing.T) {
	expected := "This is the test string"
	compressed := CompressString(expected)
	decompressed := DecompressString(compressed)
	assert.Equal(t, expected, decompressed)
	assert.True(t, len(compressed) > len(decompressed))
}

func TestCompress(t *testing.T) {
	expected := []byte("This is the test string")
	compressed := Compress(expected)
	decompressed := Decompress(compressed)
	assert.Equal(t, expected, decompressed)
	assert.True(t, len(compressed) > len(decompressed))
}
<commit_msg>Converted compress tests to test subpackage<commit_after>package util_test

import (
	"testing"

	"github.com/RobotsAndPencils/go-saml/util"
	"github.com/stretchr/testify/assert"
)

func TestCompressString(t *testing.T) {
	expected := "This is the test string"
	compressed := util.CompressString(expected)
	decompressed := util.DecompressString(compressed)
	assert.Equal(t, expected, decompressed)
	assert.True(t, len(compressed) > len(decompressed))
}

func TestCompress(t *testing.T) {
	expected := []byte("This is the test string")
	compressed := util.Compress(expected)
	decompressed := util.Decompress(compressed)
	assert.Equal(t, expected, decompressed)
	assert.True(t, len(compressed) > len(decompressed))
}
<|endoftext|>


<commit_before>package utils

import (
	"encoding/json"
	"fmt"
	"os"

	"github.com/techjanitor/pram-libs/config"
)

var Services *Capabilities

// controls capabilities for internal utils
type Capabilities struct {

	// utils
	Utils struct {
		Akismet bool
	}

	// storage capabilities
	Storage struct {
		Amazon bool
		Google bool
	}
}

func CheckServices() {

	Services = &Capabilities{}

	if config.Settings.Amazon.Key != "" {
		Services.Storage.Amazon = true
	}

	if config.Settings.Google.Key != "" {
		Services.Storage.Google = true
	}

	if config.Settings.Akismet.Key != "" {
		Services.Utils.Akismet = true
	}

}

func (c Capabilities) Print() {

	// Marshal the structs into JSON
	output, err := json.MarshalIndent(c, "", "  ")
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}

	fmt.Printf("%s\n", output)

}
<commit_msg>add local config<commit_after>package utils

import (
	"encoding/json"
	"fmt"
	"os"

	"github.com/techjanitor/pram-libs/config"
)

var Services *Capabilities

// controls capabilities for internal utils
type Capabilities struct {

	// utils
	Utils struct {
		Akismet bool
	}

	// storage capabilities
	Storage struct {
		Amazon bool
		Google bool
	}
}

func init() {
	Services = &Capabilities{}
}

func CheckServices() {

	if config.Settings.Amazon.Key != "" {
		Services.Storage.Amazon = true
	}

	if config.Settings.Google.Key != "" {
		Services.Storage.Google = true
	}

	if config.Settings.Akismet.Key != "" {
		Services.Utils.Akismet = true
	}

}

func (c Capabilities) Print() {

	// Marshal the structs into JSON
	output, err := json.MarshalIndent(c, "", "  ")
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}

	fmt.Printf("%s\n", output)

}
<|endoftext|>


<commit_before>package utils

import (
	"fmt"
	"github.com/garyburd/redigo/redis"
	"strconv"

	e "github.com/techjanitor/pram-libs/errors"
	r "github.com/techjanitor/pram-libs/redis"
)

var (
	maxLogins    int  = 5
	limitSeconds uint = 300
)

// will increment a counter in redis to limit login attempts
func LoginCounter(userid uint) (err error) {

	// convert userid to string
	uid := strconv.Itoa(int(userid))

	// Initialize cache handle
	cache := r.RedisCache

	// key is like login:21
	key := fmt.Sprintf("login:%s", uid)

	// increment login key
	result, err := cache.Incr(key)
	if err != nil {
		return e.ErrInternalError
	}

	// increment login key
	err = cache.Expire(key, limitSeconds)
	if err != nil {
		return e.ErrInternalError
	}

	if result >= maxLogins {
		return e.ErrMaxLogins
	}

	return

}
<commit_msg>add redis and audit libs<commit_after>package utils

import (
	"fmt"
	"strconv"

	e "github.com/techjanitor/pram-libs/errors"
	"github.com/techjanitor/pram-libs/redis"
)

var (
	maxLogins    int  = 5
	limitSeconds uint = 300
)

// will increment a counter in redis to limit login attempts
func LoginCounter(userid uint) (err error) {

	// convert userid to string
	uid := strconv.Itoa(int(userid))

	// Initialize cache handle
	cache := redis.RedisCache

	// key is like login:21
	key := fmt.Sprintf("login:%s", uid)

	// increment login key
	result, err := cache.Incr(key)
	if err != nil {
		return e.ErrInternalError
	}

	// increment login key
	err = cache.Expire(key, limitSeconds)
	if err != nil {
		return e.ErrInternalError
	}

	if result >= maxLogins {
		return e.ErrMaxLogins
	}

	return

}
<|endoftext|>