<commit_before>package wallet

import (
	"crypto/rand"
	"errors"
	"runtime"
	"sync"

	"github.com/NebulousLabs/Sia/crypto"
	"github.com/NebulousLabs/Sia/encoding"
	"github.com/NebulousLabs/Sia/modules"
	"github.com/NebulousLabs/Sia/types"
	"github.com/NebulousLabs/bolt"
)

var (
	errKnownSeed = errors.New("seed is already known")
)

type (
	// uniqueID is a unique id randomly generated and put at the front of every
	// persistence object. It is used to make sure that a different encryption
	// key can be used for every persistence object.
	uniqueID [crypto.EntropySize]byte

	// seedFile stores an encrypted wallet seed on disk.
	seedFile struct {
		UID                    uniqueID
		EncryptionVerification crypto.Ciphertext
		Seed                   crypto.Ciphertext
	}
)

// generateSpendableKey creates the keys and unlock conditions for seed at a
// given index.
func generateSpendableKey(seed modules.Seed, index uint64) spendableKey {
	sk, pk := crypto.GenerateKeyPairDeterministic(crypto.HashAll(seed, index))
	return spendableKey{
		UnlockConditions: types.UnlockConditions{
			PublicKeys:         []types.SiaPublicKey{types.Ed25519PublicKey(pk)},
			SignaturesRequired: 1,
		},
		SecretKeys: []crypto.SecretKey{sk},
	}
}

// generateKeys generates n keys from seed, starting from index start.
func generateKeys(seed modules.Seed, start, n uint64) []spendableKey {
	// generate in parallel, one goroutine per core.
	keys := make([]spendableKey, n)
	var wg sync.WaitGroup
	wg.Add(runtime.NumCPU())
	for cpu := 0; cpu < runtime.NumCPU(); cpu++ {
		go func(offset uint64) {
			defer wg.Done()
			for i := offset; i < n; i += uint64(runtime.NumCPU()) {
				// NOTE: don't bother trying to optimize generateSpendableKey;
				// profiling shows that ed25519 key generation consumes far
				// more CPU time than encoding or hashing.
				keys[i] = generateSpendableKey(seed, start+i)
			}
		}(uint64(cpu))
	}
	wg.Wait()
	return keys
}

// createSeedFile creates and encrypts a seedFile.
func createSeedFile(masterKey crypto.TwofishKey, seed modules.Seed) (seedFile, error) {
	var sf seedFile
	_, err := rand.Read(sf.UID[:])
	if err != nil {
		return seedFile{}, err
	}
	sek := uidEncryptionKey(masterKey, sf.UID)
	sf.EncryptionVerification, err = sek.EncryptBytes(verificationPlaintext)
	if err != nil {
		return seedFile{}, err
	}
	sf.Seed, err = sek.EncryptBytes(seed[:])
	if err != nil {
		return seedFile{}, err
	}
	return sf, nil
}

// decryptSeedFile decrypts a seed file using the encryption key.
func decryptSeedFile(masterKey crypto.TwofishKey, sf seedFile) (seed modules.Seed, err error) {
	// Verify that the provided master key is the correct key.
	decryptionKey := uidEncryptionKey(masterKey, sf.UID)
	err = verifyEncryption(decryptionKey, sf.EncryptionVerification)
	if err != nil {
		return modules.Seed{}, err
	}

	// Decrypt and return the seed.
	plainSeed, err := decryptionKey.DecryptBytes(sf.Seed)
	if err != nil {
		return modules.Seed{}, err
	}
	copy(seed[:], plainSeed)
	return seed, nil
}

// integrateSeed generates n spendableKeys from the seed and loads them into
// the wallet.
func (w *Wallet) integrateSeed(seed modules.Seed, n uint64) {
	for _, sk := range generateKeys(seed, 0, n) {
		w.keys[sk.UnlockConditions.UnlockHash()] = sk
	}
}

// loadSeed integrates a recovery seed into the wallet.
func (w *Wallet) loadSeed(masterKey crypto.TwofishKey, seed modules.Seed) error {
	// Because the recovery seed does not have a UID, duplication must be
	// prevented by comparing with the list of decrypted seeds. This can only
	// occur while the wallet is unlocked.
	if !w.unlocked {
		return modules.ErrLockedWallet
	}
	if seed == w.primarySeed {
		return errKnownSeed
	}
	for _, wSeed := range w.seeds {
		if seed == wSeed {
			return errKnownSeed
		}
	}

	err := w.db.Update(func(tx *bolt.Tx) error {
		err := checkMasterKey(tx, masterKey)
		if err != nil {
			return err
		}

		// create a seedFile for the seed
		sf, err := createSeedFile(masterKey, seed)
		if err != nil {
			return err
		}

		// add the seedFile
		return tx.Bucket(bucketSeedFiles).Put(sf.UID[:], encoding.Marshal(sf))
	})
	if err != nil {
		return err
	}

	// load the seed's keys
	w.integrateSeed(seed, modules.PublicKeysPerSeed)
	w.seeds = append(w.seeds, seed)
	return nil
}

// nextPrimarySeedAddress fetches the next address from the primary seed.
func (w *Wallet) nextPrimarySeedAddress(tx *bolt.Tx) (types.UnlockConditions, error) {
	// Check that the wallet has been unlocked.
	if !w.unlocked {
		return types.UnlockConditions{}, modules.ErrLockedWallet
	}

	// Fetch and increment the seed progress.
	progress, err := dbGetPrimarySeedProgress(tx)
	if err != nil {
		return types.UnlockConditions{}, err
	}
	if err = dbPutPrimarySeedProgress(tx, progress+1); err != nil {
		return types.UnlockConditions{}, err
	}
	// Integrate the next key into the wallet, and return the unlock
	// conditions.
	spendableKey := generateSpendableKey(w.primarySeed, progress)
	w.keys[spendableKey.UnlockConditions.UnlockHash()] = spendableKey
	return spendableKey.UnlockConditions, nil
}

// AllSeeds returns a list of all seeds known to and used by the wallet.
func (w *Wallet) AllSeeds() ([]modules.Seed, error) {
	w.mu.Lock()
	defer w.mu.Unlock()
	if !w.unlocked {
		return nil, modules.ErrLockedWallet
	}
	return append([]modules.Seed{w.primarySeed}, w.seeds...), nil
}

// PrimarySeed returns the decrypted primary seed of the wallet, as well as
// the number of addresses that the seed can be safely used to generate.
func (w *Wallet) PrimarySeed() (modules.Seed, uint64, error) {
	w.mu.Lock()
	defer w.mu.Unlock()
	if !w.unlocked {
		return modules.Seed{}, 0, modules.ErrLockedWallet
	}
	// TODO: going to the db is slow; consider caching progress. On the other
	// hand, PrimarySeed isn't a frequently called method, so caching may be
	// overkill.
	var progress uint64
	err := w.db.View(func(tx *bolt.Tx) error {
		var err error
		progress, err = dbGetPrimarySeedProgress(tx)
		return err
	})
	if err != nil {
		return modules.Seed{}, 0, err
	}

	// addresses remaining is maxScanKeys-progress; generating more keys than
	// that risks not being able to recover them when using SweepSeed or
	// InitFromSeed.
	remaining := maxScanKeys - progress
	if progress > maxScanKeys {
		remaining = 0
	}
	return w.primarySeed, remaining, nil
}

// NextAddress returns an unlock hash that is ready to receive siacoins or
// siafunds. The address is generated using the primary address seed.
func (w *Wallet) NextAddress() (types.UnlockConditions, error) {
	if err := w.tg.Add(); err != nil {
		return types.UnlockConditions{}, err
	}
	defer w.tg.Done()
	w.mu.Lock()
	defer w.mu.Unlock()

	// TODO: going to the db is slow; consider creating 100 addresses at a
	// time.
	var uc types.UnlockConditions
	err := w.db.Update(func(tx *bolt.Tx) error {
		var err error
		uc, err = w.nextPrimarySeedAddress(tx)
		return err
	})
	return uc, err
}

// LoadSeed will track all of the addresses generated by the input seed,
// reclaiming any funds that were lost due to a deleted file or lost encryption
// key. An error will be returned if the seed has already been integrated with
// the wallet.
func (w *Wallet) LoadSeed(masterKey crypto.TwofishKey, seed modules.Seed) error {
	if err := w.tg.Add(); err != nil {
		return err
	}
	defer w.tg.Done()
	w.mu.Lock()
	defer w.mu.Unlock()
	return w.loadSeed(masterKey, seed)
}

// SweepSeed scans the blockchain for outputs generated from seed and creates
// a transaction that transfers them to the wallet. Note that this incurs a
// transaction fee. It returns the total value of the outputs, minus the fee.
// If only siafunds were found, the fee is deducted from the wallet.
func (w *Wallet) SweepSeed(seed modules.Seed) (coins, funds types.Currency, err error) {
	if err = w.tg.Add(); err != nil {
		return
	}
	defer w.tg.Done()

	if !w.cs.Synced() {
		return types.Currency{}, types.Currency{}, errors.New("cannot sweep until blockchain is synced")
	}

	// get an address to spend into
	var uc types.UnlockConditions
	err = w.db.Update(func(tx *bolt.Tx) error {
		var err error
		uc, err = w.nextPrimarySeedAddress(tx)
		return err
	})
	if err != nil {
		return
	}

	// scan blockchain for outputs, filtering out 'dust' (outputs that cost
	// more in fees than they are worth)
	s := newSeedScanner(seed)
	_, maxFee := w.tpool.FeeEstimation()
	const outputSize = 350 // approx. size in bytes of an output and accompanying signature
	s.dustThreshold = maxFee.Mul64(outputSize)
	if err = s.scan(w.cs); err != nil {
		return
	}

	// construct a transaction that spends the outputs
	// TODO: this may result in transactions that are too large.
	tb := w.StartTransaction()
	var sweptCoins, sweptFunds types.Currency // total values of swept outputs
	for _, output := range s.siacoinOutputs {
		// construct a siacoin input that spends the output
		sk := generateSpendableKey(seed, output.seedIndex)
		tb.AddSiacoinInput(types.SiacoinInput{
			ParentID:         types.SiacoinOutputID(output.id),
			UnlockConditions: sk.UnlockConditions,
		})
		// add a signature for the input
		sweptCoins = sweptCoins.Add(output.value)
	}
	for _, output := range s.siafundOutputs {
		// construct a siafund input that spends the output
		sk := generateSpendableKey(seed, output.seedIndex)
		tb.AddSiafundInput(types.SiafundInput{
			ParentID:         types.SiafundOutputID(output.id),
			UnlockConditions: sk.UnlockConditions,
		})
		// add a signature for the input
		sweptFunds = sweptFunds.Add(output.value)
	}

	// estimate the transaction size and fee. NOTE: this equation doesn't
	// account for other fields in the transaction, but since we are
	// multiplying by maxFee, lowballing is ok
	estTxnSize := (len(s.siacoinOutputs) + len(s.siafundOutputs)) * outputSize
	estFee := maxFee.Mul64(uint64(estTxnSize))
	tb.AddMinerFee(estFee)

	// calculate total siacoin payout
	if sweptCoins.Cmp(estFee) > 0 {
		coins = sweptCoins.Sub(estFee)
	}
	funds = sweptFunds

	switch {
	case coins.IsZero() && funds.IsZero():
		// if we aren't sweeping any coins or funds, then just return an
		// error; no reason to proceed
		return types.Currency{}, types.Currency{}, errors.New("transaction fee exceeds value of swept outputs")

	case !coins.IsZero() && funds.IsZero():
		// if we're sweeping coins but not funds, add a siacoin output for
		// them
		tb.AddSiacoinOutput(types.SiacoinOutput{
			Value:      coins,
			UnlockHash: uc.UnlockHash(),
		})

	case coins.IsZero() && !funds.IsZero():
		// if we're sweeping funds but not coins, add a siafund output for
		// them. This is tricky because we still need to pay for the
		// transaction fee, but we can't simply subtract the fee from the
		// output value like we can with swept coins. Instead, we need to fund
		// the fee using the existing wallet balance.
		tb.AddSiafundOutput(types.SiafundOutput{
			Value:      funds,
			UnlockHash: uc.UnlockHash(),
		})
		err = tb.FundSiacoins(estFee)
		if err != nil {
			return types.Currency{}, types.Currency{}, errors.New("couldn't pay transaction fee on swept funds: " + err.Error())
		}

	case !coins.IsZero() && !funds.IsZero():
		// if we're sweeping both coins and funds, add a siacoin output and a
		// siafund output
		tb.AddSiacoinOutput(types.SiacoinOutput{
			Value:      coins,
			UnlockHash: uc.UnlockHash(),
		})
		tb.AddSiafundOutput(types.SiafundOutput{
			Value:      funds,
			UnlockHash: uc.UnlockHash(),
		})
	}

	// add signatures for all coins and funds (manually, since tb doesn't have
	// access to the signing keys)
	txn, parents := tb.View()
	for _, output := range s.siacoinOutputs {
		sk := generateSpendableKey(seed, output.seedIndex)
		addSignatures(&txn, types.FullCoveredFields, sk.UnlockConditions, crypto.Hash(output.id), sk)
	}
	for _, sfo := range s.siafundOutputs {
		sk := generateSpendableKey(seed, sfo.seedIndex)
		addSignatures(&txn, types.FullCoveredFields, sk.UnlockConditions, crypto.Hash(sfo.id), sk)
	}
	// Usually, all the inputs will come from swept outputs. However, there is
	// an edge case in which inputs will be added from the wallet. To cover
	// this case, we iterate through the SiacoinInputs and add a signature for
	// any input that belongs to the wallet.
	w.mu.RLock()
	for _, input := range txn.SiacoinInputs {
		if key, ok := w.keys[input.UnlockConditions.UnlockHash()]; ok {
			addSignatures(&txn, types.FullCoveredFields, input.UnlockConditions, crypto.Hash(input.ParentID), key)
		}
	}
	w.mu.RUnlock()

	// submit the transaction
	txnSet := append(parents, txn)
	err = w.tpool.AcceptTransactionSet(txnSet)
	return
}
<commit_msg>prevent sweeping the primary seed<commit_after>package wallet

import (
	"crypto/rand"
	"errors"
	"runtime"
	"sync"

	"github.com/NebulousLabs/Sia/crypto"
	"github.com/NebulousLabs/Sia/encoding"
	"github.com/NebulousLabs/Sia/modules"
	"github.com/NebulousLabs/Sia/types"
	"github.com/NebulousLabs/bolt"
)

var (
	errKnownSeed = errors.New("seed is already known")
)

type (
	// uniqueID is a unique id randomly generated and put at the front of every
	// persistence object. It is used to make sure that a different encryption
	// key can be used for every persistence object.
	uniqueID [crypto.EntropySize]byte

	// seedFile stores an encrypted wallet seed on disk.
	seedFile struct {
		UID                    uniqueID
		EncryptionVerification crypto.Ciphertext
		Seed                   crypto.Ciphertext
	}
)

// generateSpendableKey creates the keys and unlock conditions for seed at a
// given index.
func generateSpendableKey(seed modules.Seed, index uint64) spendableKey {
	sk, pk := crypto.GenerateKeyPairDeterministic(crypto.HashAll(seed, index))
	return spendableKey{
		UnlockConditions: types.UnlockConditions{
			PublicKeys:         []types.SiaPublicKey{types.Ed25519PublicKey(pk)},
			SignaturesRequired: 1,
		},
		SecretKeys: []crypto.SecretKey{sk},
	}
}

// generateKeys generates n keys from seed, starting from index start.
func generateKeys(seed modules.Seed, start, n uint64) []spendableKey {
	// generate in parallel, one goroutine per core.
	keys := make([]spendableKey, n)
	var wg sync.WaitGroup
	wg.Add(runtime.NumCPU())
	for cpu := 0; cpu < runtime.NumCPU(); cpu++ {
		go func(offset uint64) {
			defer wg.Done()
			for i := offset; i < n; i += uint64(runtime.NumCPU()) {
				// NOTE: don't bother trying to optimize generateSpendableKey;
				// profiling shows that ed25519 key generation consumes far
				// more CPU time than encoding or hashing.
				keys[i] = generateSpendableKey(seed, start+i)
			}
		}(uint64(cpu))
	}
	wg.Wait()
	return keys
}

// createSeedFile creates and encrypts a seedFile.
func createSeedFile(masterKey crypto.TwofishKey, seed modules.Seed) (seedFile, error) {
	var sf seedFile
	_, err := rand.Read(sf.UID[:])
	if err != nil {
		return seedFile{}, err
	}
	sek := uidEncryptionKey(masterKey, sf.UID)
	sf.EncryptionVerification, err = sek.EncryptBytes(verificationPlaintext)
	if err != nil {
		return seedFile{}, err
	}
	sf.Seed, err = sek.EncryptBytes(seed[:])
	if err != nil {
		return seedFile{}, err
	}
	return sf, nil
}

// decryptSeedFile decrypts a seed file using the encryption key.
func decryptSeedFile(masterKey crypto.TwofishKey, sf seedFile) (seed modules.Seed, err error) {
	// Verify that the provided master key is the correct key.
	decryptionKey := uidEncryptionKey(masterKey, sf.UID)
	err = verifyEncryption(decryptionKey, sf.EncryptionVerification)
	if err != nil {
		return modules.Seed{}, err
	}

	// Decrypt and return the seed.
	plainSeed, err := decryptionKey.DecryptBytes(sf.Seed)
	if err != nil {
		return modules.Seed{}, err
	}
	copy(seed[:], plainSeed)
	return seed, nil
}

// integrateSeed generates n spendableKeys from the seed and loads them into
// the wallet.
func (w *Wallet) integrateSeed(seed modules.Seed, n uint64) {
	for _, sk := range generateKeys(seed, 0, n) {
		w.keys[sk.UnlockConditions.UnlockHash()] = sk
	}
}

// loadSeed integrates a recovery seed into the wallet.
func (w *Wallet) loadSeed(masterKey crypto.TwofishKey, seed modules.Seed) error {
	// Because the recovery seed does not have a UID, duplication must be
	// prevented by comparing with the list of decrypted seeds. This can only
	// occur while the wallet is unlocked.
	if !w.unlocked {
		return modules.ErrLockedWallet
	}
	if seed == w.primarySeed {
		return errKnownSeed
	}
	for _, wSeed := range w.seeds {
		if seed == wSeed {
			return errKnownSeed
		}
	}

	err := w.db.Update(func(tx *bolt.Tx) error {
		err := checkMasterKey(tx, masterKey)
		if err != nil {
			return err
		}

		// create a seedFile for the seed
		sf, err := createSeedFile(masterKey, seed)
		if err != nil {
			return err
		}

		// add the seedFile
		return tx.Bucket(bucketSeedFiles).Put(sf.UID[:], encoding.Marshal(sf))
	})
	if err != nil {
		return err
	}

	// load the seed's keys
	w.integrateSeed(seed, modules.PublicKeysPerSeed)
	w.seeds = append(w.seeds, seed)
	return nil
}

// nextPrimarySeedAddress fetches the next address from the primary seed.
func (w *Wallet) nextPrimarySeedAddress(tx *bolt.Tx) (types.UnlockConditions, error) {
	// Check that the wallet has been unlocked.
	if !w.unlocked {
		return types.UnlockConditions{}, modules.ErrLockedWallet
	}

	// Fetch and increment the seed progress.
	progress, err := dbGetPrimarySeedProgress(tx)
	if err != nil {
		return types.UnlockConditions{}, err
	}
	if err = dbPutPrimarySeedProgress(tx, progress+1); err != nil {
		return types.UnlockConditions{}, err
	}
	// Integrate the next key into the wallet, and return the unlock
	// conditions.
	spendableKey := generateSpendableKey(w.primarySeed, progress)
	w.keys[spendableKey.UnlockConditions.UnlockHash()] = spendableKey
	return spendableKey.UnlockConditions, nil
}

// AllSeeds returns a list of all seeds known to and used by the wallet.
func (w *Wallet) AllSeeds() ([]modules.Seed, error) {
	w.mu.Lock()
	defer w.mu.Unlock()
	if !w.unlocked {
		return nil, modules.ErrLockedWallet
	}
	return append([]modules.Seed{w.primarySeed}, w.seeds...), nil
}

// PrimarySeed returns the decrypted primary seed of the wallet, as well as
// the number of addresses that the seed can be safely used to generate.
func (w *Wallet) PrimarySeed() (modules.Seed, uint64, error) {
	w.mu.Lock()
	defer w.mu.Unlock()
	if !w.unlocked {
		return modules.Seed{}, 0, modules.ErrLockedWallet
	}
	// TODO: going to the db is slow; consider caching progress. On the other
	// hand, PrimarySeed isn't a frequently called method, so caching may be
	// overkill.
	var progress uint64
	err := w.db.View(func(tx *bolt.Tx) error {
		var err error
		progress, err = dbGetPrimarySeedProgress(tx)
		return err
	})
	if err != nil {
		return modules.Seed{}, 0, err
	}

	// addresses remaining is maxScanKeys-progress; generating more keys than
	// that risks not being able to recover them when using SweepSeed or
	// InitFromSeed.
	remaining := maxScanKeys - progress
	if progress > maxScanKeys {
		remaining = 0
	}
	return w.primarySeed, remaining, nil
}

// NextAddress returns an unlock hash that is ready to receive siacoins or
// siafunds. The address is generated using the primary address seed.
func (w *Wallet) NextAddress() (types.UnlockConditions, error) {
	if err := w.tg.Add(); err != nil {
		return types.UnlockConditions{}, err
	}
	defer w.tg.Done()
	w.mu.Lock()
	defer w.mu.Unlock()

	// TODO: going to the db is slow; consider creating 100 addresses at a
	// time.
	var uc types.UnlockConditions
	err := w.db.Update(func(tx *bolt.Tx) error {
		var err error
		uc, err = w.nextPrimarySeedAddress(tx)
		return err
	})
	return uc, err
}

// LoadSeed will track all of the addresses generated by the input seed,
// reclaiming any funds that were lost due to a deleted file or lost encryption
// key. An error will be returned if the seed has already been integrated with
// the wallet.
func (w *Wallet) LoadSeed(masterKey crypto.TwofishKey, seed modules.Seed) error {
	if err := w.tg.Add(); err != nil {
		return err
	}
	defer w.tg.Done()
	w.mu.Lock()
	defer w.mu.Unlock()
	return w.loadSeed(masterKey, seed)
}

// SweepSeed scans the blockchain for outputs generated from seed and creates
// a transaction that transfers them to the wallet. Note that this incurs a
// transaction fee. It returns the total value of the outputs, minus the fee.
// If only siafunds were found, the fee is deducted from the wallet.
func (w *Wallet) SweepSeed(seed modules.Seed) (coins, funds types.Currency, err error) {
	if err = w.tg.Add(); err != nil {
		return
	}
	defer w.tg.Done()

	w.mu.RLock()
	match := seed == w.primarySeed
	w.mu.RUnlock()
	if match {
		return types.Currency{}, types.Currency{}, errors.New("cannot sweep primary seed")
	}

	if !w.cs.Synced() {
		return types.Currency{}, types.Currency{}, errors.New("cannot sweep until blockchain is synced")
	}

	// get an address to spend into
	var uc types.UnlockConditions
	err = w.db.Update(func(tx *bolt.Tx) error {
		var err error
		uc, err = w.nextPrimarySeedAddress(tx)
		return err
	})
	if err != nil {
		return
	}

	// scan blockchain for outputs, filtering out 'dust' (outputs that cost
	// more in fees than they are worth)
	s := newSeedScanner(seed)
	_, maxFee := w.tpool.FeeEstimation()
	const outputSize = 350 // approx. size in bytes of an output and accompanying signature
	s.dustThreshold = maxFee.Mul64(outputSize)
	if err = s.scan(w.cs); err != nil {
		return
	}

	// construct a transaction that spends the outputs
	// TODO: this may result in transactions that are too large.
	tb := w.StartTransaction()
	var sweptCoins, sweptFunds types.Currency // total values of swept outputs
	for _, output := range s.siacoinOutputs {
		// construct a siacoin input that spends the output
		sk := generateSpendableKey(seed, output.seedIndex)
		tb.AddSiacoinInput(types.SiacoinInput{
			ParentID:         types.SiacoinOutputID(output.id),
			UnlockConditions: sk.UnlockConditions,
		})
		// add a signature for the input
		sweptCoins = sweptCoins.Add(output.value)
	}
	for _, output := range s.siafundOutputs {
		// construct a siafund input that spends the output
		sk := generateSpendableKey(seed, output.seedIndex)
		tb.AddSiafundInput(types.SiafundInput{
			ParentID:         types.SiafundOutputID(output.id),
			UnlockConditions: sk.UnlockConditions,
		})
		// add a signature for the input
		sweptFunds = sweptFunds.Add(output.value)
	}

	// estimate the transaction size and fee. NOTE: this equation doesn't
	// account for other fields in the transaction, but since we are
	// multiplying by maxFee, lowballing is ok
	estTxnSize := (len(s.siacoinOutputs) + len(s.siafundOutputs)) * outputSize
	estFee := maxFee.Mul64(uint64(estTxnSize))
	tb.AddMinerFee(estFee)

	// calculate total siacoin payout
	if sweptCoins.Cmp(estFee) > 0 {
		coins = sweptCoins.Sub(estFee)
	}
	funds = sweptFunds

	switch {
	case coins.IsZero() && funds.IsZero():
		// if we aren't sweeping any coins or funds, then just return an
		// error; no reason to proceed
		return types.Currency{}, types.Currency{}, errors.New("transaction fee exceeds value of swept outputs")

	case !coins.IsZero() && funds.IsZero():
		// if we're sweeping coins but not funds, add a siacoin output for
		// them
		tb.AddSiacoinOutput(types.SiacoinOutput{
			Value:      coins,
			UnlockHash: uc.UnlockHash(),
		})

	case coins.IsZero() && !funds.IsZero():
		// if we're sweeping funds but not coins, add a siafund output for
		// them. This is tricky because we still need to pay for the
		// transaction fee, but we can't simply subtract the fee from the
		// output value like we can with swept coins. Instead, we need to fund
		// the fee using the existing wallet balance.
		tb.AddSiafundOutput(types.SiafundOutput{
			Value:      funds,
			UnlockHash: uc.UnlockHash(),
		})
		err = tb.FundSiacoins(estFee)
		if err != nil {
			return types.Currency{}, types.Currency{}, errors.New("couldn't pay transaction fee on swept funds: " + err.Error())
		}

	case !coins.IsZero() && !funds.IsZero():
		// if we're sweeping both coins and funds, add a siacoin output and a
		// siafund output
		tb.AddSiacoinOutput(types.SiacoinOutput{
			Value:      coins,
			UnlockHash: uc.UnlockHash(),
		})
		tb.AddSiafundOutput(types.SiafundOutput{
			Value:      funds,
			UnlockHash: uc.UnlockHash(),
		})
	}

	// add signatures for all coins and funds (manually, since tb doesn't have
	// access to the signing keys)
	txn, parents := tb.View()
	for _, output := range s.siacoinOutputs {
		sk := generateSpendableKey(seed, output.seedIndex)
		addSignatures(&txn, types.FullCoveredFields, sk.UnlockConditions, crypto.Hash(output.id), sk)
	}
	for _, sfo := range s.siafundOutputs {
		sk := generateSpendableKey(seed, sfo.seedIndex)
		addSignatures(&txn, types.FullCoveredFields, sk.UnlockConditions, crypto.Hash(sfo.id), sk)
	}
	// Usually, all the inputs will come from swept outputs. However, there is
	// an edge case in which inputs will be added from the wallet. To cover
	// this case, we iterate through the SiacoinInputs and add a signature for
	// any input that belongs to the wallet.
	w.mu.RLock()
	for _, input := range txn.SiacoinInputs {
		if key, ok := w.keys[input.UnlockConditions.UnlockHash()]; ok {
			addSignatures(&txn, types.FullCoveredFields, input.UnlockConditions, crypto.Hash(input.ParentID), key)
		}
	}
	w.mu.RUnlock()

	// submit the transaction
	txnSet := append(parents, txn)
	err = w.tpool.AcceptTransactionSet(txnSet)
	return
}
<|endoftext|>


<commit_before>package torrent

import (
	"net"
	"net/http"
	"net/url"
	"time"

	"github.com/anacrolix/dht/v2"
	"github.com/anacrolix/dht/v2/krpc"
	"github.com/anacrolix/log"
	"github.com/anacrolix/missinggo/v2"
	"github.com/anacrolix/missinggo/v2/expect"
	"github.com/anacrolix/torrent/version"
	"golang.org/x/time/rate"

	"github.com/anacrolix/torrent/iplist"
	"github.com/anacrolix/torrent/mse"
	"github.com/anacrolix/torrent/storage"
)

// Probably not safe to modify this after it's given to a Client.
type ClientConfig struct {
	// Store torrent file data in this directory unless .DefaultStorage is
	// specified.
	DataDir string `long:"data-dir" description:"directory to store downloaded torrent data"`
	// The address to listen for new uTP and TCP BitTorrent protocol connections. DHT shares a UDP
	// socket with uTP unless configured otherwise.
	ListenHost              func(network string) string
	ListenPort              int
	NoDefaultPortForwarding bool
	UpnpID                  string
	// Don't announce to trackers. This only leaves DHT to discover peers.
	DisableTrackers bool `long:"disable-trackers"`
	DisablePEX      bool `long:"disable-pex"`

	// Don't create a DHT.
	NoDHT            bool `long:"disable-dht"`
	DhtStartingNodes func(network string) dht.StartingNodesGetter
	// Called for each anacrolix/dht Server created for the Client.
	ConfigureAnacrolixDhtServer       func(*dht.ServerConfig)
	PeriodicallyAnnounceTorrentsToDht bool

	// Never send chunks to peers.
	NoUpload bool `long:"no-upload"`
	// Disable uploading even when it isn't fair.
	DisableAggressiveUpload bool `long:"disable-aggressive-upload"`
	// Upload even after there's nothing in it for us. By default uploading is
	// not altruistic, we'll only upload to encourage the peer to reciprocate.
	Seed bool `long:"seed"`
	// Only applies to chunks uploaded to peers, to maintain responsiveness
	// communicating local Client state to peers. Each limiter token
	// represents one byte. The Limiter's burst must be large enough to fit a
	// whole chunk, which is usually 16 KiB (see TorrentSpec.ChunkSize).
	UploadRateLimiter *rate.Limiter
	// Rate limits all reads from connections to peers. Each limiter token
	// represents one byte. The Limiter's burst must be bigger than the
	// largest Read performed on a the underlying rate-limiting io.Reader
	// minus one. This is likely to be the larger of the main read loop buffer
	// (~4096), and the requested chunk size (~16KiB, see
	// TorrentSpec.ChunkSize).
	DownloadRateLimiter *rate.Limiter
	// Maximum unverified bytes across all torrents. Not used if zero.
	MaxUnverifiedBytes int64

	// User-provided Client peer ID. If not present, one is generated automatically.
	PeerID string
	// For the bittorrent protocol.
	DisableUTP bool
	// For the bittorrent protocol.
	DisableTCP bool `long:"disable-tcp"`
	// Called to instantiate storage for each added torrent. Builtin backends
	// are in the storage package. If not set, the "file" implementation is
	// used (and Closed when the Client is Closed).
	DefaultStorage storage.ClientImpl

	HeaderObfuscationPolicy HeaderObfuscationPolicy
	// The crypto methods to offer when initiating connections with header obfuscation.
	CryptoProvides mse.CryptoMethod
	// Chooses the crypto method to use when receiving connections with header obfuscation.
	CryptoSelector mse.CryptoSelector

	IPBlocklist      iplist.Ranger
	DisableIPv6      bool `long:"disable-ipv6"`
	DisableIPv4      bool
	DisableIPv4Peers bool
	// Perform logging and any other behaviour that will help debug.
	Debug  bool `help:"enable debugging"`
	Logger log.Logger

	// Defines proxy for HTTP requests, such as for trackers. It's commonly set from the result of
	// "net/http".ProxyURL(HTTPProxy).
	HTTPProxy func(*http.Request) (*url.URL, error)
	// HTTPUserAgent changes default UserAgent for HTTP requests
	HTTPUserAgent string
	// Updated occasionally to when there's been some changes to client
	// behaviour in case other clients are assuming anything of us. See also
	// `bep20`.
	ExtendedHandshakeClientVersion string
	// Peer ID client identifier prefix. We'll update this occasionally to
	// reflect changes to client behaviour that other clients may depend on.
	// Also see `extendedHandshakeClientVersion`.
	Bep20 string

	// Peer dial timeout to use when there are limited peers.
	NominalDialTimeout time.Duration
	// Minimum peer dial timeout to use (even if we have lots of peers).
	MinDialTimeout             time.Duration
	EstablishedConnsPerTorrent int
	HalfOpenConnsPerTorrent    int
	TotalHalfOpenConns         int
	// Maximum number of peer addresses in reserve.
	TorrentPeersHighWater int
	// Minumum number of peers before effort is made to obtain more peers.
	TorrentPeersLowWater int

	// Limit how long handshake can take. This is to reduce the lingering
	// impact of a few bad apples. 4s loses 1% of successful handshakes that
	// are obtained with 60s timeout, and 5% of unsuccessful handshakes.
	HandshakesTimeout time.Duration
	// How long between writes before sending a keep alive message on a peer connection that we want
	// to maintain.
	KeepAliveTimeout time.Duration

	// The IP addresses as our peers should see them. May differ from the
	// local interfaces due to NAT or other network configurations.
	PublicIp4 net.IP
	PublicIp6 net.IP

	// Accept rate limiting affects excessive connection attempts from IPs that fail during
	// handshakes or request torrents that we don't have.
	DisableAcceptRateLimiting bool
	// Don't add connections that have the same peer ID as an existing
	// connection for a given Torrent.
	DropDuplicatePeerIds bool
	// Drop peers that are complete if we are also complete and have no use for the peer. This is a
	// bit of a special case, since a peer could also be useless if they're just not interested, or
	// we don't intend to obtain all of a torrent's data.
	DropMutuallyCompletePeers bool
	// Whether to accept peer connections at all.
	AcceptPeerConnections bool
	// Whether a Client should want conns without delegating to any attached Torrents. This is
	// useful when torrents might be added dynmically in callbacks for example.
	AlwaysWantConns bool

	// OnQuery hook func
	DHTOnQuery func(query *krpc.Msg, source net.Addr) (propagate bool)

	Extensions PeerExtensionBits

	DisableWebtorrent bool
	DisableWebseeds   bool

	Callbacks Callbacks
}

func (cfg *ClientConfig) SetListenAddr(addr string) *ClientConfig {
	host, port, err := missinggo.ParseHostPort(addr)
	expect.Nil(err)
	cfg.ListenHost = func(string) string { return host }
	cfg.ListenPort = port
	return cfg
}

func NewDefaultClientConfig() *ClientConfig {
	cc := &ClientConfig{
		HTTPUserAgent:                  version.DefaultHttpUserAgent,
		ExtendedHandshakeClientVersion: version.DefaultExtendedHandshakeClientVersion,
		Bep20:                          version.DefaultBep20Prefix,
		UpnpID:                         version.DefaultUpnpId,
		NominalDialTimeout:             20 * time.Second,
		MinDialTimeout:                 3 * time.Second,
		EstablishedConnsPerTorrent:     50,
		HalfOpenConnsPerTorrent:        25,
		TotalHalfOpenConns:             100,
		TorrentPeersHighWater:          500,
		TorrentPeersLowWater:           50,
		HandshakesTimeout:              4 * time.Second,
		KeepAliveTimeout:               time.Minute,
		DhtStartingNodes: func(network string) dht.StartingNodesGetter {
			return func() ([]dht.Addr, error) { return dht.GlobalBootstrapAddrs(network) }
		},
		PeriodicallyAnnounceTorrentsToDht: true,
		ListenHost:                        func(string) string { return "" },
		UploadRateLimiter:                 unlimited,
		DownloadRateLimiter:               unlimited,
		DisableAcceptRateLimiting:         true,
		DropMutuallyCompletePeers:         true,
		HeaderObfuscationPolicy: HeaderObfuscationPolicy{
			Preferred:        true,
			RequirePreferred: false,
		},
		CryptoSelector:        mse.DefaultCryptoSelector,
		CryptoProvides:        mse.AllSupportedCrypto,
		ListenPort:            42069,
		Extensions:            defaultPeerExtensionBytes(),
		AcceptPeerConnections: true,
	}
	//cc.ConnTracker.SetNoMaxEntries()
	//cc.ConnTracker.Timeout = func(conntrack.Entry) time.Duration { return 0 }
	return cc
}

type HeaderObfuscationPolicy struct {
	RequirePreferred bool // Whether the value of Preferred is a strict requirement.
	Preferred        bool // Whether header obfuscation is preferred.
}
<commit_msg>Eliminate `expect` dependency (#667)<commit_after>package torrent

import (
	"net"
	"net/http"
	"net/url"
	"time"

	"github.com/anacrolix/dht/v2"
	"github.com/anacrolix/dht/v2/krpc"
	"github.com/anacrolix/log"
	"github.com/anacrolix/missinggo/v2"
	"github.com/anacrolix/torrent/version"
	"golang.org/x/time/rate"

	"github.com/anacrolix/torrent/iplist"
	"github.com/anacrolix/torrent/mse"
	"github.com/anacrolix/torrent/storage"
)

// Probably not safe to modify this after it's given to a Client.
type ClientConfig struct {
	// Store torrent file data in this directory unless .DefaultStorage is
	// specified.
	DataDir string `long:"data-dir" description:"directory to store downloaded torrent data"`
	// The address to listen for new uTP and TCP BitTorrent protocol connections. DHT shares a UDP
	// socket with uTP unless configured otherwise.
	ListenHost              func(network string) string
	ListenPort              int
	NoDefaultPortForwarding bool
	UpnpID                  string
	// Don't announce to trackers. This only leaves DHT to discover peers.
	DisableTrackers bool `long:"disable-trackers"`
	DisablePEX      bool `long:"disable-pex"`

	// Don't create a DHT.
	NoDHT            bool `long:"disable-dht"`
	DhtStartingNodes func(network string) dht.StartingNodesGetter
	// Called for each anacrolix/dht Server created for the Client.
	ConfigureAnacrolixDhtServer       func(*dht.ServerConfig)
	PeriodicallyAnnounceTorrentsToDht bool

	// Never send chunks to peers.
	NoUpload bool `long:"no-upload"`
	// Disable uploading even when it isn't fair.
	DisableAggressiveUpload bool `long:"disable-aggressive-upload"`
	// Upload even after there's nothing in it for us. By default uploading is
	// not altruistic, we'll only upload to encourage the peer to reciprocate.
	Seed bool `long:"seed"`
	// Only applies to chunks uploaded to peers, to maintain responsiveness
	// communicating local Client state to peers. Each limiter token
	// represents one byte. The Limiter's burst must be large enough to fit a
	// whole chunk, which is usually 16 KiB (see TorrentSpec.ChunkSize).
	UploadRateLimiter *rate.Limiter
	// Rate limits all reads from connections to peers. Each limiter token
	// represents one byte. The Limiter's burst must be bigger than the
	// largest Read performed on a the underlying rate-limiting io.Reader
	// minus one. This is likely to be the larger of the main read loop buffer
	// (~4096), and the requested chunk size (~16KiB, see
	// TorrentSpec.ChunkSize).
	DownloadRateLimiter *rate.Limiter
	// Maximum unverified bytes across all torrents. Not used if zero.
	MaxUnverifiedBytes int64

	// User-provided Client peer ID. If not present, one is generated automatically.
	PeerID string
	// For the bittorrent protocol.
	DisableUTP bool
	// For the bittorrent protocol.
	DisableTCP bool `long:"disable-tcp"`
	// Called to instantiate storage for each added torrent. Builtin backends
	// are in the storage package. If not set, the "file" implementation is
	// used (and Closed when the Client is Closed).
	DefaultStorage storage.ClientImpl

	HeaderObfuscationPolicy HeaderObfuscationPolicy
	// The crypto methods to offer when initiating connections with header obfuscation.
	CryptoProvides mse.CryptoMethod
	// Chooses the crypto method to use when receiving connections with header obfuscation.
	CryptoSelector mse.CryptoSelector

	IPBlocklist      iplist.Ranger
	DisableIPv6      bool `long:"disable-ipv6"`
	DisableIPv4      bool
	DisableIPv4Peers bool
	// Perform logging and any other behaviour that will help debug.
	Debug  bool `help:"enable debugging"`
	Logger log.Logger

	// Defines proxy for HTTP requests, such as for trackers. It's commonly set from the result of
	// "net/http".ProxyURL(HTTPProxy).
	HTTPProxy func(*http.Request) (*url.URL, error)
	// HTTPUserAgent changes default UserAgent for HTTP requests
	HTTPUserAgent string
	// Updated occasionally to when there's been some changes to client
	// behaviour in case other clients are assuming anything of us. See also
	// `bep20`.
	ExtendedHandshakeClientVersion string
	// Peer ID client identifier prefix. We'll update this occasionally to
	// reflect changes to client behaviour that other clients may depend on.
	// Also see `extendedHandshakeClientVersion`.
	Bep20 string

	// Peer dial timeout to use when there are limited peers.
	NominalDialTimeout time.Duration
	// Minimum peer dial timeout to use (even if we have lots of peers).
	MinDialTimeout             time.Duration
	EstablishedConnsPerTorrent int
	HalfOpenConnsPerTorrent    int
	TotalHalfOpenConns         int
	// Maximum number of peer addresses in reserve.
	TorrentPeersHighWater int
	// Minumum number of peers before effort is made to obtain more peers.
	TorrentPeersLowWater int

	// Limit how long handshake can take. This is to reduce the lingering
	// impact of a few bad apples. 4s loses 1% of successful handshakes that
	// are obtained with 60s timeout, and 5% of unsuccessful handshakes.
	HandshakesTimeout time.Duration
	// How long between writes before sending a keep alive message on a peer connection that we want
	// to maintain.
	KeepAliveTimeout time.Duration

	// The IP addresses as our peers should see them. May differ from the
	// local interfaces due to NAT or other network configurations.
	PublicIp4 net.IP
	PublicIp6 net.IP

	// Accept rate limiting affects excessive connection attempts from IPs that fail during
	// handshakes or request torrents that we don't have.
	DisableAcceptRateLimiting bool
	// Don't add connections that have the same peer ID as an existing
	// connection for a given Torrent.
	DropDuplicatePeerIds bool
	// Drop peers that are complete if we are also complete and have no use for the peer. This is a
	// bit of a special case, since a peer could also be useless if they're just not interested, or
	// we don't intend to obtain all of a torrent's data.
	DropMutuallyCompletePeers bool
	// Whether to accept peer connections at all.
	AcceptPeerConnections bool
	// Whether a Client should want conns without delegating to any attached Torrents. This is
	// useful when torrents might be added dynmically in callbacks for example.
	AlwaysWantConns bool

	// OnQuery hook func
	DHTOnQuery func(query *krpc.Msg, source net.Addr) (propagate bool)

	Extensions PeerExtensionBits

	DisableWebtorrent bool
	DisableWebseeds   bool

	Callbacks Callbacks
}

func (cfg *ClientConfig) SetListenAddr(addr string) *ClientConfig {
	host, port, err := missinggo.ParseHostPort(addr)
	if err != nil {
		panic(err)
	}
	cfg.ListenHost = func(string) string { return host }
	cfg.ListenPort = port
	return cfg
}

func NewDefaultClientConfig() *ClientConfig {
	cc := &ClientConfig{
		HTTPUserAgent:                  version.DefaultHttpUserAgent,
		ExtendedHandshakeClientVersion: version.DefaultExtendedHandshakeClientVersion,
		Bep20:                          version.DefaultBep20Prefix,
		UpnpID:                         version.DefaultUpnpId,
		NominalDialTimeout:             20 * time.Second,
		MinDialTimeout:                 3 * time.Second,
		EstablishedConnsPerTorrent:     50,
		HalfOpenConnsPerTorrent:        25,
		TotalHalfOpenConns:             100,
		TorrentPeersHighWater:          500,
		TorrentPeersLowWater:           50,
		HandshakesTimeout:              4 * time.Second,
		KeepAliveTimeout:               time.Minute,
		DhtStartingNodes: func(network string) dht.StartingNodesGetter {
			return func() ([]dht.Addr, error) { return dht.GlobalBootstrapAddrs(network) }
		},
		PeriodicallyAnnounceTorrentsToDht: true,
		ListenHost:                        func(string) string { return "" },
		UploadRateLimiter:                 unlimited,
		DownloadRateLimiter:               unlimited,
		DisableAcceptRateLimiting:         true,
		DropMutuallyCompletePeers:         true,
		HeaderObfuscationPolicy: HeaderObfuscationPolicy{
			Preferred:        true,
			RequirePreferred: false,
		},
		CryptoSelector:        mse.DefaultCryptoSelector,
		CryptoProvides:        mse.AllSupportedCrypto,
		ListenPort:            42069,
		Extensions:            defaultPeerExtensionBytes(),
		AcceptPeerConnections: true,
	}
	//cc.ConnTracker.SetNoMaxEntries()
	//cc.ConnTracker.Timeout = func(conntrack.Entry) time.Duration { return 0 }
	return cc
}

type HeaderObfuscationPolicy struct {
	RequirePreferred bool // Whether the value of Preferred is a strict requirement.
	Preferred        bool // Whether header obfuscation is preferred.
}
<|endoftext|>


<commit_before>// Permission is hereby granted, free of charge, to any person
// obtaining a copy of this software and associated documentation
// files (the "Software"), to deal in the Software without
// restriction, including without limitation the rights to use,
// copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the
// Software is furnished to do so, subject to the following
// conditions:
// The above copyright notice and this permission notice shall be
// included in all copies or substantial portions of the Software.
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
// EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES
// OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
// NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
// HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
// WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
// OTHER DEALINGS IN THE SOFTWARE.
//
// See http://formwork-io.github.io/ for more.

package main

import (
	"errors"
	"fmt"
	toml "github.com/BurntSushi/toml"
	"io/ioutil"
	"os"
	"regexp"
	"strconv"
	"strings"
)

type rail struct {
	Name    string
	Pattern string
	Ingress int
	Egress  int
}

type rails struct {
	Rail []rail
}

func ReadConfigFile(path string) ([]rail, error) {
	data, err := ioutil.ReadFile(path)
	if err != nil {
		return nil, errors.New(fmt.Sprintf("configuration file error: %s", err.Error()))
	}

	var rails rails
	_, err = toml.Decode(string(data), &rails)
	if err != nil {
		return nil, errors.New(fmt.Sprintf("configuration file error: %s", err.Error()))
	}

	_, err = validateRails(rails.Rail)
	if err != nil {
		return nil, errors.New(fmt.Sprintf("configuration file error: " + err.Error()))
	}

	return rails.Rail, nil
}

func ReadEnvironment() ([]rail, error) {
	GL_RAIL_NAME_TMPL := "GL_RAIL_%d_NAME"
	GL_RAIL_PATTERN_TMPL := "GL_RAIL_%d_PATTERN"
	GL_RAIL_INGRESS_TMPL := "GL_RAIL_%d_INGRESS_PORT"
	GL_RAIL_EGRESS_TMPL := "GL_RAIL_%d_EGRESS_PORT"

	var rails []rail
	index := 0
	for {
		name, err := getenv(fmt.Sprintf(GL_RAIL_NAME_TMPL, index))
		if err != nil {
			return nil, err
		}
		if name == "" {
			break
		}

		pattern, err := getenv(fmt.Sprintf(GL_RAIL_PATTERN_TMPL, index))
		if err != nil {
			return nil, err
		}
		ingress, err := getenv(fmt.Sprintf(GL_RAIL_INGRESS_TMPL, index))
		if err != nil {
			return nil, err
		}
		egress, err := getenv(fmt.Sprintf(GL_RAIL_EGRESS_TMPL, index))
		if err != nil {
			return nil, err
		}
		ingress_port, err := asPort(ingress)
		if err != nil {
			return nil, err
		}
		egress_port, err := asPort(egress)
		if err != nil {
			return nil, err
		}

		rails = append(rails, rail{
			Name:    name,
			Pattern: pattern,
			Ingress: ingress_port,
			Egress:  egress_port,
		})

		index++
	}

	_, err := validateRails(rails)
	if err != nil {
		return nil, errors.New(fmt.Sprintf("configuration file error: " + err.Error()))
	}

	return rails, nil
}

func validateRails(rails []rail) (*rail, error) {
	re := regexp.MustCompile("[[:punct:]]|[[:space:]]")
	for i, rail := range rails {
		normalizedPattern := strings.ToLower(re.ReplaceAllLiteralString(rail.Pattern, ""))
		if normalizedPattern != "pubsub" && normalizedPattern != "reqrep" {
			msg := fmt.Sprintf("invalid pattern \"%s\" for rail \"%s\"", rail.Pattern, rail.Name)
			return &rail, errors.New(msg)
		}
		rails[i].Pattern = normalizedPattern
	}

	return nil, nil
}

func getenv(env string) (string, error) {
	_env := os.Getenv(env)
	if len(_env) == 0 {
		return "", errors.New(fmt.Sprintf("no %s is set", env))
	}
	return _env, nil
}

func asPort(env string) (int, error) {
	port, err := strconv.Atoi(env)
	if err != nil {
		die("invalid port: %s", env)
		return -1, errors.New(fmt.Sprintf("invalid port: %v - %s", env, err.Error()))
	} else if port < 1 || port > 65535 {
		die("invalid port: %s", env)
		return -1, errors.New(fmt.Sprintf("invalid port: %v - %s", env, err.Error()))
	}
	return port, nil
}
<commit_msg>break reading env if at least one rail is present<commit_after>// Permission is hereby granted, free of charge, to any person
// obtaining a copy of this software and associated documentation
// files (the "Software"), to deal in the Software without
// restriction, including without limitation the rights to use,
// copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the
// Software is furnished to do so, subject to the following
// conditions:
// The above copyright notice and this permission notice shall be
// included in all copies or substantial portions of the Software.
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
// EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES
// OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
// NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
// HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
// WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
// OTHER DEALINGS IN THE SOFTWARE.
//
// See http://formwork-io.github.io/ for more.

package main

import (
	"errors"
	"fmt"
	toml "github.com/BurntSushi/toml"
	"io/ioutil"
	"os"
	"regexp"
	"strconv"
	"strings"
)

type rail struct {
	Name    string
	Pattern string
	Ingress int
	Egress  int
}

type rails struct {
	Rail []rail
}

func ReadConfigFile(path string) ([]rail, error) {
	data, err := ioutil.ReadFile(path)
	if err != nil {
		return nil, errors.New(fmt.Sprintf("configuration file error: %s", err.Error()))
	}

	var rails rails
	_, err = toml.Decode(string(data), &rails)
	if err != nil {
		return nil, errors.New(fmt.Sprintf("configuration file error: %s", err.Error()))
	}

	_, err = validateRails(rails.Rail)
	if err != nil {
		return nil, errors.New(fmt.Sprintf("configuration file error: " + err.Error()))
	}

	return rails.Rail, nil
}

func ReadEnvironment() ([]rail, error) {
	GL_RAIL_NAME_TMPL := "GL_RAIL_%d_NAME"
	GL_RAIL_PATTERN_TMPL := "GL_RAIL_%d_PATTERN"
	GL_RAIL_INGRESS_TMPL := "GL_RAIL_%d_INGRESS_PORT"
	GL_RAIL_EGRESS_TMPL := "GL_RAIL_%d_EGRESS_PORT"

	var rails []rail
	index := 0
	for {
		name, err := getenv(fmt.Sprintf(GL_RAIL_NAME_TMPL, index))
		if err != nil {
			if index != 0 {
				break
			}
		}
		if name == "" {
			break
		}

		pattern, err := getenv(fmt.Sprintf(GL_RAIL_PATTERN_TMPL, index))
		if err != nil {
			return nil, err
		}
		ingress, err := getenv(fmt.Sprintf(GL_RAIL_INGRESS_TMPL, index))
		if err != nil {
			return nil, err
		}
		egress, err := getenv(fmt.Sprintf(GL_RAIL_EGRESS_TMPL, index))
		if err != nil {
			return nil, err
		}
		ingress_port, err := asPort(ingress)
		if err != nil {
			return nil, err
		}
		egress_port, err := asPort(egress)
		if err != nil {
			return nil, err
		}

		rails = append(rails, rail{
			Name:    name,
			Pattern: pattern,
			Ingress: ingress_port,
			Egress:  egress_port,
		})

		index++
	}

	_, err := validateRails(rails)
	if err != nil {
		return nil, errors.New(fmt.Sprintf("configuration file error: " + err.Error()))
	}

	return rails, nil
}

func validateRails(rails []rail) (*rail, error) {
	re := regexp.MustCompile("[[:punct:]]|[[:space:]]")
	for i, rail := range rails {
		normalizedPattern := strings.ToLower(re.ReplaceAllLiteralString(rail.Pattern, ""))
		if normalizedPattern != "pubsub" && normalizedPattern != "reqrep" {
			msg := fmt.Sprintf("invalid pattern \"%s\" for rail \"%s\"", rail.Pattern, rail.Name)
			return &rail, errors.New(msg)
		}
		rails[i].Pattern = normalizedPattern
	}

	return nil, nil
}

func getenv(env string) (string, error) {
	_env := os.Getenv(env)
	if len(_env) == 0 {
		return "", errors.New(fmt.Sprintf("no %s is set", env))
	}
	return _env, nil
}

func asPort(env string) (int, error) {
	port, err := strconv.Atoi(env)
	if err != nil {
		die("invalid port: %s", env)
		return -1, errors.New(fmt.Sprintf("invalid port: %v - %s", env, err.Error()))
	} else if port < 1 || port > 65535 {
		die("invalid port: %s", env)
		return -1, errors.New(fmt.Sprintf("invalid port: %v - %s", env, err.Error()))
	}
	return port, nil
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"io/ioutil"
)

type CheckConf struct {
	Handler     string
	Command     string
	Interval    int
	Standalone  bool
	Subscribers []string
  Occurrences int
}

type ClientConf struct {
	Name          string
	Address       string
	Subscriptions []string
}

type RabbitmqConf struct {
	Port     int
	Host     string
	User     string
	Password string
	Vhost    string
}

func ParseClientConfig(file string) (ClientConf, error) {
	clientJson, err := ioutil.ReadFile(file)
	if err != nil {
		return ClientConf{}, err
	}

	var clientConfig map[string]ClientConf

	if err := json.Unmarshal(clientJson, &clientConfig); err != nil {
		return ClientConf{}, err
	}
	return clientConfig["client"], nil
}

func ParseRabbitmqConfig(file string) (RabbitmqConf, error) {
	rabbitmqJson, err := ioutil.ReadFile(file)
	if err != nil {
		return RabbitmqConf{}, err
	}

	var rabbitmqConfig map[string]RabbitmqConf

	if err := json.Unmarshal(rabbitmqJson, &rabbitmqConfig); err != nil {
		return RabbitmqConf{}, err
	}
	return rabbitmqConfig["rabbitmq"], nil
}

func ParseChecksConfig(file string) (map[string]CheckConf, error) {
	checksJson, err := ioutil.ReadFile(file)
	if err != nil {
		return nil, err
	}

	var checksConfig map[string]map[string]CheckConf

	if err := json.Unmarshal(checksJson, &checksConfig); err != nil {
		return nil, err
	}
	return checksConfig["checks"], nil
}
<commit_msg>Added json tags to structs<commit_after>package main

import (
	"encoding/json"
	"io/ioutil"
)

type CheckConf struct {
	Handler     string    `json:"handler"`
	Command     string    `json:"command"`
	Interval    int       `json:"interval"`
	Standalone  bool      `json:"standalone"`
	Subscribers []string  `json:"subscribers"`
  Occurrences int       `json:"occurrences"`
}

type ClientConf struct {
	Name          string    `json:"name"`
	Address       string    `json:"address"`
	Subscriptions []string  `json:"subscriptions"`
}

type RabbitmqConf struct {
	Port     int    `json:"port"`
	Host     string `json:"host"`
	User     string `json:"user"`
	Password string `json:"password"`
	Vhost    string `json:"vhost"`
}

func ParseClientConfig(file string) (ClientConf, error) {
	clientJson, err := ioutil.ReadFile(file)
	if err != nil {
		return ClientConf{}, err
	}

	var clientConfig map[string]ClientConf

	if err := json.Unmarshal(clientJson, &clientConfig); err != nil {
		return ClientConf{}, err
	}
	return clientConfig["client"], nil
}

func ParseRabbitmqConfig(file string) (RabbitmqConf, error) {
	rabbitmqJson, err := ioutil.ReadFile(file)
	if err != nil {
		return RabbitmqConf{}, err
	}

	var rabbitmqConfig map[string]RabbitmqConf

	if err := json.Unmarshal(rabbitmqJson, &rabbitmqConfig); err != nil {
		return RabbitmqConf{}, err
	}
	return rabbitmqConfig["rabbitmq"], nil
}

func ParseChecksConfig(file string) (map[string]CheckConf, error) {
	checksJson, err := ioutil.ReadFile(file)
	if err != nil {
		return nil, err
	}

	var checksConfig map[string]map[string]CheckConf

	if err := json.Unmarshal(checksJson, &checksConfig); err != nil {
		return nil, err
	}
	return checksConfig["checks"], nil
}
<|endoftext|>


<commit_before>package goConfig

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"reflect"

	"github.com/pkg/errors"
)

const (
	defaultPath       = "./"
	defaultConfigFile = "config.json"
)

// Load config file
func Load(config interface{}) (err error) {
	configFile := defaultPath + defaultConfigFile
	file, err := os.Open(configFile)
	if err != nil {
		return
	}
	defer file.Close()

	decoder := json.NewDecoder(file)
	err = decoder.Decode(&config)
	if err != nil {
		return
	}
	return
}

// Save config file
func Save(config interface{}) (err error) {
	_, err = os.Stat(defaultPath)

	if os.IsNotExist(err) {
		os.Mkdir(defaultPath, 0700)
	}

	configFile := defaultPath + defaultConfigFile

	_, err = os.Stat(configFile)
	if err != nil {
		return
	}

	b, err := json.MarshalIndent(config, "", "\t")
	if err != nil {
		return
	}

	err = ioutil.WriteFile(defaultConfigFile, b, 0644)
	if err != nil {
		return
	}
	return
}

// Getenv get enviroment variable
func Getenv(env string) (r string) {
	r = os.Getenv(env)
	return
}

func parseTags(s interface{}) (err error) {

	st := reflect.TypeOf(s)
	vt := reflect.ValueOf(s)

	if st.Kind() != reflect.Ptr {
		err = errors.New("Not a pointer")
		return
	}

	ref := st.Elem()
	if ref.Kind() != reflect.Struct {
		err = errors.New("Not a struct")
		return
	}

	for i := 0; i < ref.NumField(); i++ {
		field := ref.Field(i)

		kindStr := ""
		if field.Type.Kind() == reflect.Struct {
			kindStr = "Struct"

			value := vt.Elem().Field(i)
			err = parseTags(value.Addr().Interface())
			if err != nil {
				return
			}

		}

		// change value POC
		if field.Type.Kind() == reflect.String {
			value := vt.Elem().Field(i)
			value.SetString("TEST")
		}

		fmt.Println("name:", field.Name,
			"| cfg:", field.Tag.Get("config"),
			"| cfgDefault:", field.Tag.Get("cfgDefault"),
			"| type:", field.Type, kindStr)

	}
	return
}
<commit_msg>reduce redundant code<commit_after>package goConfig

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"reflect"

	"github.com/pkg/errors"
)

const (
	defaultPath       = "./"
	defaultConfigFile = "config.json"
)

// Load config file
func Load(config interface{}) (err error) {
	configFile := defaultPath + defaultConfigFile
	file, err := os.Open(configFile)
	if err != nil {
		return
	}
	defer file.Close()

	decoder := json.NewDecoder(file)
	err = decoder.Decode(&config)
	if err != nil {
		return
	}
	return
}

// Save config file
func Save(config interface{}) (err error) {
	_, err = os.Stat(defaultPath)

	if os.IsNotExist(err) {
		os.Mkdir(defaultPath, 0700)
	}

	configFile := defaultPath + defaultConfigFile

	_, err = os.Stat(configFile)
	if err != nil {
		return
	}

	b, err := json.MarshalIndent(config, "", "\t")
	if err != nil {
		return
	}

	err = ioutil.WriteFile(defaultConfigFile, b, 0644)
	if err != nil {
		return
	}
	return
}

// Getenv get enviroment variable
func Getenv(env string) (r string) {
	r = os.Getenv(env)
	return
}

func parseTags(s interface{}) (err error) {

	st := reflect.TypeOf(s)
	vt := reflect.ValueOf(s)

	if st.Kind() != reflect.Ptr {
		err = errors.New("Not a pointer")
		return
	}

	ref := st.Elem()
	if ref.Kind() != reflect.Struct {
		err = errors.New("Not a struct")
		return
	}

	for i := 0; i < ref.NumField(); i++ {
		field := ref.Field(i)
		value := vt.Elem().Field(i)

		kindStr := ""
		if field.Type.Kind() == reflect.Struct {
			kindStr = "Struct"

			err = parseTags(value.Addr().Interface())
			if err != nil {
				return
			}

		}

		// change value POC
		if field.Type.Kind() == reflect.String {
			value.SetString("TEST")
		}

		fmt.Println("name:", field.Name,
			"| cfg:", field.Tag.Get("config"),
			"| cfgDefault:", field.Tag.Get("cfgDefault"),
			"| type:", field.Type, kindStr)

	}
	return
}
<|endoftext|>


<commit_before>package faas

import (
	"io/ioutil"
	"os"
	"path/filepath"

	"gopkg.in/yaml.v2"
)

// ConfigFile is the name of the config's serialized form.
const ConfigFile = ".faas.yaml"

// Config represents the serialized state of a Function's metadata.
// See the Function struct for attribute documentation.
type config struct {
	Name      string `yaml:"name"`
	Namespace string `yaml:"namespace"`
	Runtime   string `yaml:"runtime"`
	Image     string `yaml:"image"`
	// Add new values to the toConfig/fromConfig functions.
}

// newConfig returns a Config populated from data serialized to disk if it is
// available.  Errors are returned if the path is not valid, if there are
// errors accessing an extant config file, or the contents of the file do not
// unmarshall.  A missing file at a valid path does not error but returns the
// empty value of Config.
func newConfig(root string) (c config, err error) {
	filename := filepath.Join(root, ConfigFile)
	if _, err = os.Stat(filename); os.IsNotExist(err) {
		err = nil // do not consider a missing config file an error
		return    // return the zero value of the config
	}
	bb, err := ioutil.ReadFile(filename)
	if err != nil {
		return
	}
	err = yaml.Unmarshal(bb, &c)
	return
}

// fromConfig returns a Function populated from config.
// Note that config does not include ancillary fields not serialized, such as Root.
func fromConfig(c config) (f Function) {
	return Function{
		Name:      c.Name,
		Namespace: c.Namespace,
		Runtime:   c.Runtime,
		Image:     c.Image,
	}
}

// toConfig serializes a Function to a config object.
func toConfig(f Function) config {
	return config{
		Name:      f.Name,
		Namespace: f.Namespace,
		Runtime:   f.Runtime,
		Image:     f.Image,
	}
}

// writeConfig for the given Function out to disk at root.
func writeConfig(f Function) (err error) {
	path := filepath.Join(f.Root, ConfigFile)
	c := toConfig(f)
	var bb []byte
	if bb, err = yaml.Marshal(&c); err != nil {
		return
	}
	return ioutil.WriteFile(path, bb, 0644)
}
<commit_msg>fix: serialize trigger on faas.config<commit_after>package faas

import (
	"io/ioutil"
	"os"
	"path/filepath"

	"gopkg.in/yaml.v2"
)

// ConfigFile is the name of the config's serialized form.
const ConfigFile = ".faas.yaml"

// Config represents the serialized state of a Function's metadata.
// See the Function struct for attribute documentation.
type config struct {
	Name      string `yaml:"name"`
	Namespace string `yaml:"namespace"`
	Runtime   string `yaml:"runtime"`
	Image     string `yaml:"image"`
	Trigger   string `yaml:"trigger"`
	// Add new values to the toConfig/fromConfig functions.
}

// newConfig returns a Config populated from data serialized to disk if it is
// available.  Errors are returned if the path is not valid, if there are
// errors accessing an extant config file, or the contents of the file do not
// unmarshall.  A missing file at a valid path does not error but returns the
// empty value of Config.
func newConfig(root string) (c config, err error) {
	filename := filepath.Join(root, ConfigFile)
	if _, err = os.Stat(filename); os.IsNotExist(err) {
		err = nil // do not consider a missing config file an error
		return    // return the zero value of the config
	}
	bb, err := ioutil.ReadFile(filename)
	if err != nil {
		return
	}
	err = yaml.Unmarshal(bb, &c)
	return
}

// fromConfig returns a Function populated from config.
// Note that config does not include ancillary fields not serialized, such as Root.
func fromConfig(c config) (f Function) {
	return Function{
		Name:      c.Name,
		Namespace: c.Namespace,
		Runtime:   c.Runtime,
		Image:     c.Image,
		Trigger:   c.Trigger,
	}
}

// toConfig serializes a Function to a config object.
func toConfig(f Function) config {
	return config{
		Name:      f.Name,
		Namespace: f.Namespace,
		Runtime:   f.Runtime,
		Image:     f.Image,
		Trigger:   f.Trigger,
	}
}

// writeConfig for the given Function out to disk at root.
func writeConfig(f Function) (err error) {
	path := filepath.Join(f.Root, ConfigFile)
	c := toConfig(f)
	var bb []byte
	if bb, err = yaml.Marshal(&c); err != nil {
		return
	}
	return ioutil.WriteFile(path, bb, 0644)
}
<|endoftext|>


<commit_before>/*
Copyright 2018 The Kubernetes Authors.

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

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

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

package gce

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"time"

	"github.com/golang/glog"
	"google.golang.org/api/googleapi"
	tpuapi "google.golang.org/api/tpu/v1"

	"k8s.io/apimachinery/pkg/util/wait"
)

// newTPUService returns a new tpuService using the client to communicate with
// the Cloud TPU APIs.
func newTPUService(client *http.Client) (*tpuService, error) {
	s, err := tpuapi.New(client)
	if err != nil {
		return nil, err
	}
	return &tpuService{
		nodesService:      tpuapi.NewProjectsLocationsNodesService(s),
		operationsService: tpuapi.NewProjectsLocationsOperationsService(s),
	}, nil
}

// tpuService encapsulates the TPU services on nodes and the operations on the
// nodes.
type tpuService struct {
	nodesService      *tpuapi.ProjectsLocationsNodesService
	operationsService *tpuapi.ProjectsLocationsOperationsService
}

// CreateTPU creates the Cloud TPU node with the specified name in the
// specified zone.
func (gce *GCECloud) CreateTPU(ctx context.Context, name, zone string, node *tpuapi.Node) (*tpuapi.Node, error) {
	var err error
	mc := newTPUMetricContext("create", zone)
	defer mc.Observe(err)

	var op *tpuapi.Operation
	parent := getTPUParentName(gce.projectID, zone)
	op, err = gce.tpuService.nodesService.Create(parent, node).NodeId(name).Do()
	if err != nil {
		return nil, err
	}
	glog.V(2).Infof("Creating Cloud TPU %q in zone %q with operation %q", name, zone, op.Name)

	op, err = gce.waitForTPUOp(ctx, op)
	if err != nil {
		return nil, err
	}
	err = getErrorFromTPUOp(op)
	if err != nil {
		return nil, err
	}

	output := new(tpuapi.Node)
	err = json.Unmarshal(op.Response, output)
	if err != nil {
		err = fmt.Errorf("failed to unmarshal response from operation %q: response = %v, err = %v", op.Name, op.Response, err)
		return nil, err
	}
	return output, nil
}

// DeleteTPU deletes the Cloud TPU with the specified name in the specified
// zone.
func (gce *GCECloud) DeleteTPU(ctx context.Context, name, zone string) error {
	var err error
	mc := newTPUMetricContext("delete", zone)
	defer mc.Observe(err)

	var op *tpuapi.Operation
	name = getTPUName(gce.projectID, zone, name)
	op, err = gce.tpuService.nodesService.Delete(name).Do()
	if err != nil {
		return err
	}
	glog.V(2).Infof("Deleting Cloud TPU %q in zone %q with operation %q", name, zone, op.Name)

	op, err = gce.waitForTPUOp(ctx, op)
	if err != nil {
		return err
	}
	err = getErrorFromTPUOp(op)
	if err != nil {
		return err
	}
	return nil
}

// GetTPU returns the Cloud TPU with the specified name in the specified zone.
func (gce *GCECloud) GetTPU(ctx context.Context, name, zone string) (*tpuapi.Node, error) {
	mc := newTPUMetricContext("get", zone)

	name = getTPUName(gce.projectID, zone, name)
	node, err := gce.tpuService.nodesService.Get(name).Do()
	if err != nil {
		return nil, mc.Observe(err)
	}
	return node, mc.Observe(nil)
}

// ListTPUs returns Cloud TPUs in the specified zone.
func (gce *GCECloud) ListTPUs(ctx context.Context, zone string) ([]*tpuapi.Node, error) {
	mc := newTPUMetricContext("list", zone)

	parent := getTPUParentName(gce.projectID, zone)
	response, err := gce.tpuService.nodesService.List(parent).Do()
	if err != nil {
		return nil, mc.Observe(err)
	}
	return response.Nodes, mc.Observe(nil)
}

// waitForTPUOp checks whether the op is done every 30 seconds before the ctx
// is cancelled.
func (gce *GCECloud) waitForTPUOp(ctx context.Context, op *tpuapi.Operation) (*tpuapi.Operation, error) {
	if err := wait.PollInfinite(30*time.Second, func() (bool, error) {
		// Check if context has been cancelled.
		select {
		case <-ctx.Done():
			glog.V(3).Infof("Context for operation %q has been cancelled: %s", op.Name, ctx.Err())
			return true, ctx.Err()
		default:
		}

		glog.V(3).Infof("Waiting for operation %q to complete...", op.Name)

		start := time.Now()
		gce.operationPollRateLimiter.Accept()
		duration := time.Now().Sub(start)
		if duration > 5*time.Second {
			glog.V(2).Infof("Getting operation %q throttled for %v", op.Name, duration)
		}

		var err error
		op, err = gce.tpuService.operationsService.Get(op.Name).Do()
		if err != nil {
			return true, err
		}
		if op.Done {
			glog.V(3).Infof("Operation %q has completed", op.Name)
			return true, nil
		}
		return false, nil
	}); err != nil {
		return nil, fmt.Errorf("failed to wait for operation %q: %s", op.Name, err)
	}
	return op, nil
}

// newTPUMetricContext returns a new metricContext used for recording metrics
// of Cloud TPU API calls.
func newTPUMetricContext(request, zone string) *metricContext {
	return newGenericMetricContext("tpus", request, unusedMetricLabel, zone, "v1")
}

// getErrorFromTPUOp returns the error in the failed op, or nil if the op
// succeed.
func getErrorFromTPUOp(op *tpuapi.Operation) error {
	if op != nil && op.Error != nil {
		return &googleapi.Error{
			Code:    op.ServerResponse.HTTPStatusCode,
			Message: op.Error.Message,
		}
	}
	return nil
}

func getTPUParentName(project, zone string) string {
	return fmt.Sprintf("projects/%s/locations/%s", project, zone)
}

func getTPUName(project, zone, name string) string {
	return fmt.Sprintf("projects/%s/locations/%s/nodes/%s", project, zone, name)
}
<commit_msg>GCE: Add ListLocations to Cloud TPU API<commit_after>/*
Copyright 2018 The Kubernetes Authors.

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

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

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

package gce

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"time"

	"github.com/golang/glog"
	"google.golang.org/api/googleapi"
	tpuapi "google.golang.org/api/tpu/v1"

	"k8s.io/apimachinery/pkg/util/wait"
)

// newTPUService returns a new tpuService using the client to communicate with
// the Cloud TPU APIs.
func newTPUService(client *http.Client) (*tpuService, error) {
	s, err := tpuapi.New(client)
	if err != nil {
		return nil, err
	}
	return &tpuService{
		projects: tpuapi.NewProjectsService(s),
	}, nil
}

// tpuService encapsulates the TPU services on nodes and the operations on the
// nodes.
type tpuService struct {
	projects *tpuapi.ProjectsService
}

// CreateTPU creates the Cloud TPU node with the specified name in the
// specified zone.
func (gce *GCECloud) CreateTPU(ctx context.Context, name, zone string, node *tpuapi.Node) (*tpuapi.Node, error) {
	var err error
	mc := newTPUMetricContext("create", zone)
	defer mc.Observe(err)

	var op *tpuapi.Operation
	parent := getTPUParentName(gce.projectID, zone)
	op, err = gce.tpuService.projects.Locations.Nodes.Create(parent, node).NodeId(name).Do()
	if err != nil {
		return nil, err
	}
	glog.V(2).Infof("Creating Cloud TPU %q in zone %q with operation %q", name, zone, op.Name)

	op, err = gce.waitForTPUOp(ctx, op)
	if err != nil {
		return nil, err
	}
	err = getErrorFromTPUOp(op)
	if err != nil {
		return nil, err
	}

	output := new(tpuapi.Node)
	err = json.Unmarshal(op.Response, output)
	if err != nil {
		err = fmt.Errorf("failed to unmarshal response from operation %q: response = %v, err = %v", op.Name, op.Response, err)
		return nil, err
	}
	return output, nil
}

// DeleteTPU deletes the Cloud TPU with the specified name in the specified
// zone.
func (gce *GCECloud) DeleteTPU(ctx context.Context, name, zone string) error {
	var err error
	mc := newTPUMetricContext("delete", zone)
	defer mc.Observe(err)

	var op *tpuapi.Operation
	name = getTPUName(gce.projectID, zone, name)
	op, err = gce.tpuService.projects.Locations.Nodes.Delete(name).Do()
	if err != nil {
		return err
	}
	glog.V(2).Infof("Deleting Cloud TPU %q in zone %q with operation %q", name, zone, op.Name)

	op, err = gce.waitForTPUOp(ctx, op)
	if err != nil {
		return err
	}
	err = getErrorFromTPUOp(op)
	if err != nil {
		return err
	}
	return nil
}

// GetTPU returns the Cloud TPU with the specified name in the specified zone.
func (gce *GCECloud) GetTPU(ctx context.Context, name, zone string) (*tpuapi.Node, error) {
	mc := newTPUMetricContext("get", zone)

	name = getTPUName(gce.projectID, zone, name)
	node, err := gce.tpuService.projects.Locations.Nodes.Get(name).Do()
	if err != nil {
		return nil, mc.Observe(err)
	}
	return node, mc.Observe(nil)
}

// ListTPUs returns Cloud TPUs in the specified zone.
func (gce *GCECloud) ListTPUs(ctx context.Context, zone string) ([]*tpuapi.Node, error) {
	mc := newTPUMetricContext("list", zone)

	parent := getTPUParentName(gce.projectID, zone)
	response, err := gce.tpuService.projects.Locations.Nodes.List(parent).Do()
	if err != nil {
		return nil, mc.Observe(err)
	}
	return response.Nodes, mc.Observe(nil)
}

// ListLocations returns the zones where Cloud TPUs are available.
func (gce *GCECloud) ListLocations(ctx context.Context) ([]*tpuapi.Location, error) {
	mc := newTPUMetricContext("list_locations", "")
	parent := getTPUProjectURL(gce.projectID)
	response, err := gce.tpuService.projects.Locations.List(parent).Do()
	if err != nil {
		return nil, mc.Observe(err)
	}
	return response.Locations, mc.Observe(nil)
}

// waitForTPUOp checks whether the op is done every 30 seconds before the ctx
// is cancelled.
func (gce *GCECloud) waitForTPUOp(ctx context.Context, op *tpuapi.Operation) (*tpuapi.Operation, error) {
	if err := wait.PollInfinite(30*time.Second, func() (bool, error) {
		// Check if context has been cancelled.
		select {
		case <-ctx.Done():
			glog.V(3).Infof("Context for operation %q has been cancelled: %s", op.Name, ctx.Err())
			return true, ctx.Err()
		default:
		}

		glog.V(3).Infof("Waiting for operation %q to complete...", op.Name)

		start := time.Now()
		gce.operationPollRateLimiter.Accept()
		duration := time.Now().Sub(start)
		if duration > 5*time.Second {
			glog.V(2).Infof("Getting operation %q throttled for %v", op.Name, duration)
		}

		var err error
		op, err = gce.tpuService.projects.Locations.Operations.Get(op.Name).Do()
		if err != nil {
			return true, err
		}
		if op.Done {
			glog.V(3).Infof("Operation %q has completed", op.Name)
			return true, nil
		}
		return false, nil
	}); err != nil {
		return nil, fmt.Errorf("failed to wait for operation %q: %s", op.Name, err)
	}
	return op, nil
}

// newTPUMetricContext returns a new metricContext used for recording metrics
// of Cloud TPU API calls.
func newTPUMetricContext(request, zone string) *metricContext {
	return newGenericMetricContext("tpus", request, unusedMetricLabel, zone, "v1")
}

// getErrorFromTPUOp returns the error in the failed op, or nil if the op
// succeed.
func getErrorFromTPUOp(op *tpuapi.Operation) error {
	if op != nil && op.Error != nil {
		return &googleapi.Error{
			Code:    op.ServerResponse.HTTPStatusCode,
			Message: op.Error.Message,
		}
	}
	return nil
}

func getTPUProjectURL(project string) string {
	return fmt.Sprintf("projects/%s", project)
}

func getTPUParentName(project, zone string) string {
	return fmt.Sprintf("projects/%s/locations/%s", project, zone)
}

func getTPUName(project, zone, name string) string {
	return fmt.Sprintf("projects/%s/locations/%s/nodes/%s", project, zone, name)
}
<|endoftext|>


<commit_before>/*
Copyright 2016 The Kubernetes Authors.

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

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

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

package metricsutil

import (
	"fmt"
	"io"
	"sort"

	"k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/resource"
	"k8s.io/kubernetes/pkg/api/legacyscheme"
	"k8s.io/kubernetes/pkg/printers"
	metricsapi "k8s.io/metrics/pkg/apis/metrics"
)

var (
	MeasuredResources = []v1.ResourceName{
		v1.ResourceCPU,
		v1.ResourceMemory,
	}
	NodeColumns     = []string{"NAME", "CPU(cores)", "CPU%", "MEMORY(bytes)", "MEMORY%"}
	PodColumns      = []string{"NAME", "CPU(cores)", "MEMORY(bytes)"}
	NamespaceColumn = "NAMESPACE"
	PodColumn       = "POD"
)

type ResourceMetricsInfo struct {
	Name      string
	Metrics   v1.ResourceList
	Available v1.ResourceList
}

type TopCmdPrinter struct {
	out io.Writer
}

func NewTopCmdPrinter(out io.Writer) *TopCmdPrinter {
	return &TopCmdPrinter{out: out}
}

func (printer *TopCmdPrinter) PrintNodeMetrics(metrics []metricsapi.NodeMetrics, availableResources map[string]v1.ResourceList) error {
	if len(metrics) == 0 {
		return nil
	}
	w := printers.GetNewTabWriter(printer.out)
	defer w.Flush()

	sort.Slice(metrics, func(i, j int) bool {
		return metrics[i].Name < metrics[j].Name
	})

	printColumnNames(w, NodeColumns)
	var usage v1.ResourceList
	for _, m := range metrics {
		err := legacyscheme.Scheme.Convert(&m.Usage, &usage, nil)
		if err != nil {
			return err
		}
		printMetricsLine(w, &ResourceMetricsInfo{
			Name:      m.Name,
			Metrics:   usage,
			Available: availableResources[m.Name],
		})
	}
	return nil
}

func (printer *TopCmdPrinter) PrintPodMetrics(metrics []metricsapi.PodMetrics, printContainers bool, withNamespace bool) error {
	if len(metrics) == 0 {
		return nil
	}
	w := printers.GetNewTabWriter(printer.out)
	defer w.Flush()

	if withNamespace {
		printValue(w, NamespaceColumn)
	}
	if printContainers {
		printValue(w, PodColumn)
	}

	sort.Slice(metrics, func(i, j int) bool {
		if withNamespace && metrics[i].Namespace != metrics[j].Namespace {
			return metrics[i].Namespace < metrics[j].Namespace
		}
		return metrics[i].Name < metrics[j].Name
	})

	printColumnNames(w, PodColumns)
	for _, m := range metrics {
		err := printSinglePodMetrics(w, &m, printContainers, withNamespace)
		if err != nil {
			return err
		}
	}
	return nil
}

func printColumnNames(out io.Writer, names []string) {
	for _, name := range names {
		printValue(out, name)
	}
	fmt.Fprint(out, "\n")
}

func printSinglePodMetrics(out io.Writer, m *metricsapi.PodMetrics, printContainersOnly bool, withNamespace bool) error {
	containers := make(map[string]v1.ResourceList)
	podMetrics := make(v1.ResourceList)
	for _, res := range MeasuredResources {
		podMetrics[res], _ = resource.ParseQuantity("0")
	}

	for _, c := range m.Containers {
		var usage v1.ResourceList
		err := legacyscheme.Scheme.Convert(&c.Usage, &usage, nil)
		if err != nil {
			return err
		}
		containers[c.Name] = usage
		if !printContainersOnly {
			for _, res := range MeasuredResources {
				quantity := podMetrics[res]
				quantity.Add(usage[res])
				podMetrics[res] = quantity
			}
		}
	}
	if printContainersOnly {
		for contName := range containers {
			if withNamespace {
				printValue(out, m.Namespace)
			}
			printValue(out, m.Name)
			printMetricsLine(out, &ResourceMetricsInfo{
				Name:      contName,
				Metrics:   containers[contName],
				Available: v1.ResourceList{},
			})
		}
	} else {
		if withNamespace {
			printValue(out, m.Namespace)
		}
		printMetricsLine(out, &ResourceMetricsInfo{
			Name:      m.Name,
			Metrics:   podMetrics,
			Available: v1.ResourceList{},
		})
	}
	return nil
}

func printMetricsLine(out io.Writer, metrics *ResourceMetricsInfo) {
	printValue(out, metrics.Name)
	printAllResourceUsages(out, metrics)
	fmt.Fprint(out, "\n")
}

func printValue(out io.Writer, value interface{}) {
	fmt.Fprintf(out, "%v\t", value)
}

func printAllResourceUsages(out io.Writer, metrics *ResourceMetricsInfo) {
	for _, res := range MeasuredResources {
		quantity := metrics.Metrics[res]
		printSingleResourceUsage(out, res, quantity)
		fmt.Fprint(out, "\t")
		if available, found := metrics.Available[res]; found {
			fraction := float64(quantity.MilliValue()) / float64(available.MilliValue()) * 100
			fmt.Fprintf(out, "%d%%\t", int64(fraction))
		}
	}
}

func printSingleResourceUsage(out io.Writer, resourceType v1.ResourceName, quantity resource.Quantity) {
	switch resourceType {
	case v1.ResourceCPU:
		fmt.Fprintf(out, "%vm", quantity.MilliValue())
	case v1.ResourceMemory:
		fmt.Fprintf(out, "%vMi", quantity.Value()/(1024*1024))
	default:
		fmt.Fprintf(out, "%v", quantity.Value())
	}
}
<commit_msg>print nodes for those metrics is somehow unreachable<commit_after>/*
Copyright 2016 The Kubernetes Authors.

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

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

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

package metricsutil

import (
	"fmt"
	"io"
	"sort"

	"k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/resource"
	"k8s.io/kubernetes/pkg/api/legacyscheme"
	"k8s.io/kubernetes/pkg/printers"
	metricsapi "k8s.io/metrics/pkg/apis/metrics"
)

var (
	MeasuredResources = []v1.ResourceName{
		v1.ResourceCPU,
		v1.ResourceMemory,
	}
	NodeColumns     = []string{"NAME", "CPU(cores)", "CPU%", "MEMORY(bytes)", "MEMORY%"}
	PodColumns      = []string{"NAME", "CPU(cores)", "MEMORY(bytes)"}
	NamespaceColumn = "NAMESPACE"
	PodColumn       = "POD"
)

type ResourceMetricsInfo struct {
	Name      string
	Metrics   v1.ResourceList
	Available v1.ResourceList
}

type TopCmdPrinter struct {
	out io.Writer
}

func NewTopCmdPrinter(out io.Writer) *TopCmdPrinter {
	return &TopCmdPrinter{out: out}
}

func (printer *TopCmdPrinter) PrintNodeMetrics(metrics []metricsapi.NodeMetrics, availableResources map[string]v1.ResourceList) error {
	if len(metrics) == 0 {
		return nil
	}
	w := printers.GetNewTabWriter(printer.out)
	defer w.Flush()

	sort.Slice(metrics, func(i, j int) bool {
		return metrics[i].Name < metrics[j].Name
	})

	printColumnNames(w, NodeColumns)
	var usage v1.ResourceList
	for _, m := range metrics {
		err := legacyscheme.Scheme.Convert(&m.Usage, &usage, nil)
		if err != nil {
			return err
		}
		printMetricsLine(w, &ResourceMetricsInfo{
			Name:      m.Name,
			Metrics:   usage,
			Available: availableResources[m.Name],
		})
		delete(availableResources, m.Name)
	}

	// print lines for nodes of which the metrics is unreachable.
	for nodeName := range availableResources {
		printMissingMetricsNodeLine(w, nodeName)
	}
	return nil
}

func (printer *TopCmdPrinter) PrintPodMetrics(metrics []metricsapi.PodMetrics, printContainers bool, withNamespace bool) error {
	if len(metrics) == 0 {
		return nil
	}
	w := printers.GetNewTabWriter(printer.out)
	defer w.Flush()

	if withNamespace {
		printValue(w, NamespaceColumn)
	}
	if printContainers {
		printValue(w, PodColumn)
	}

	sort.Slice(metrics, func(i, j int) bool {
		if withNamespace && metrics[i].Namespace != metrics[j].Namespace {
			return metrics[i].Namespace < metrics[j].Namespace
		}
		return metrics[i].Name < metrics[j].Name
	})

	printColumnNames(w, PodColumns)
	for _, m := range metrics {
		err := printSinglePodMetrics(w, &m, printContainers, withNamespace)
		if err != nil {
			return err
		}
	}
	return nil
}

func printColumnNames(out io.Writer, names []string) {
	for _, name := range names {
		printValue(out, name)
	}
	fmt.Fprint(out, "\n")
}

func printSinglePodMetrics(out io.Writer, m *metricsapi.PodMetrics, printContainersOnly bool, withNamespace bool) error {
	containers := make(map[string]v1.ResourceList)
	podMetrics := make(v1.ResourceList)
	for _, res := range MeasuredResources {
		podMetrics[res], _ = resource.ParseQuantity("0")
	}

	for _, c := range m.Containers {
		var usage v1.ResourceList
		err := legacyscheme.Scheme.Convert(&c.Usage, &usage, nil)
		if err != nil {
			return err
		}
		containers[c.Name] = usage
		if !printContainersOnly {
			for _, res := range MeasuredResources {
				quantity := podMetrics[res]
				quantity.Add(usage[res])
				podMetrics[res] = quantity
			}
		}
	}
	if printContainersOnly {
		for contName := range containers {
			if withNamespace {
				printValue(out, m.Namespace)
			}
			printValue(out, m.Name)
			printMetricsLine(out, &ResourceMetricsInfo{
				Name:      contName,
				Metrics:   containers[contName],
				Available: v1.ResourceList{},
			})
		}
	} else {
		if withNamespace {
			printValue(out, m.Namespace)
		}
		printMetricsLine(out, &ResourceMetricsInfo{
			Name:      m.Name,
			Metrics:   podMetrics,
			Available: v1.ResourceList{},
		})
	}
	return nil
}

func printMetricsLine(out io.Writer, metrics *ResourceMetricsInfo) {
	printValue(out, metrics.Name)
	printAllResourceUsages(out, metrics)
	fmt.Fprint(out, "\n")
}

func printMissingMetricsNodeLine(out io.Writer, nodeName string) {
	printValue(out, nodeName)
	for i := 0; i < len(MeasuredResources); i++ {
		printValue(out, "-\t-\t")
	}
	fmt.Fprint(out, "\n")
}

func printValue(out io.Writer, value interface{}) {
	fmt.Fprintf(out, "%v\t", value)
}

func printAllResourceUsages(out io.Writer, metrics *ResourceMetricsInfo) {
	for _, res := range MeasuredResources {
		quantity := metrics.Metrics[res]
		printSingleResourceUsage(out, res, quantity)
		fmt.Fprint(out, "\t")
		if available, found := metrics.Available[res]; found {
			fraction := float64(quantity.MilliValue()) / float64(available.MilliValue()) * 100
			fmt.Fprintf(out, "%d%%\t", int64(fraction))
		}
	}
}

func printSingleResourceUsage(out io.Writer, resourceType v1.ResourceName, quantity resource.Quantity) {
	switch resourceType {
	case v1.ResourceCPU:
		fmt.Fprintf(out, "%vm", quantity.MilliValue())
	case v1.ResourceMemory:
		fmt.Fprintf(out, "%vMi", quantity.Value()/(1024*1024))
	default:
		fmt.Fprintf(out, "%v", quantity.Value())
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io"
	"log"
	"os"
	"syscall"
	"unsafe"
)

/*** defines ***/

const KILO_VERSION = "0.0.1"

/*** data ***/

type Termios struct {
	Iflag  uint32
	Oflag  uint32
	Cflag  uint32
	Lflag  uint32
	Cc     [20]byte
	Ispeed uint32
	Ospeed uint32
}

type editorConfig struct {
	cx          int
	cy          int
	screenRows  int
	screenCols  int
	origTermios *Termios
}

type WinSize struct {
	Row    uint16
	Col    uint16
	Xpixel uint16
	Ypixel uint16
}

var E editorConfig

/*** terminal ***/

func die(err error) {
	disableRawMode()
	io.WriteString(os.Stdout, "\x1b[2J")
	io.WriteString(os.Stdout, "\x1b[H")
	log.Fatal(err)
}

func TcSetAttr(fd uintptr, termios *Termios) error {
	// TCSETS+1 == TCSETSW, because TCSAFLUSH doesn't exist
	if _, _, err := syscall.Syscall(syscall.SYS_IOCTL, fd, uintptr(syscall.TCSETS+1), uintptr(unsafe.Pointer(termios))); err != 0 {
		return err
	}
	return nil
}

func TcGetAttr(fd uintptr) *Termios {
	var termios = &Termios{}
	if _, _, err := syscall.Syscall(syscall.SYS_IOCTL, fd, syscall.TCGETS, uintptr(unsafe.Pointer(termios))); err != 0 {
		log.Fatalf("Problem getting terminal attributes: %s\n", err)
	}
	return termios
}

func enableRawMode() {
	E.origTermios = TcGetAttr(os.Stdin.Fd())
	var raw Termios
	raw = *E.origTermios
	raw.Iflag &^= syscall.BRKINT | syscall.ICRNL | syscall.INPCK | syscall.ISTRIP | syscall.IXON
	raw.Oflag &^= syscall.OPOST
	raw.Cflag |= syscall.CS8
	raw.Lflag &^= syscall.ECHO | syscall.ICANON | syscall.IEXTEN | syscall.ISIG
	raw.Cc[syscall.VMIN+1] = 0
	raw.Cc[syscall.VTIME+1] = 1
	if e := TcSetAttr(os.Stdin.Fd(), &raw); e != nil {
		log.Fatalf("Problem enabling raw mode: %s\n", e)
	}
}

func disableRawMode() {
	if e := TcSetAttr(os.Stdin.Fd(), E.origTermios); e != nil {
		log.Fatalf("Problem disabling raw mode: %s\n", e)
	}
}

func editorReadKey() byte {
	var buffer [1]byte
	var cc int
	var err error
	for cc, err = os.Stdin.Read(buffer[:]); cc != 1; cc, err = os.Stdin.Read(buffer[:]) {
	}
	if err != nil {
		die(err)
	}
	return buffer[0]
}

func getCursorPosition(rows *int, cols *int) int {
	io.WriteString(os.Stdout, "\x1b[6n")
	var buffer [1]byte
	var buf []byte
	var cc int
	for cc, _ = os.Stdin.Read(buffer[:]); cc == 1; cc, _ = os.Stdin.Read(buffer[:]) {
		if buffer[0] == 'R' {
			break
		}
		buf = append(buf, buffer[0])
	}
	if string(buf[0:2]) != "\x1b[" {
		log.Printf("Failed to read rows;cols from tty\n")
		return -1
	}
	if n, e := fmt.Sscanf(string(buf[2:]), "%d;%d", rows, cols); n != 2 || e != nil {
		if e != nil {
			log.Printf("getCursorPosition: fmt.Sscanf() failed: %s\n", e)
		}
		if n != 2 {
			log.Printf("getCursorPosition: got %d items, wanted 2\n", n)
		}
		return -1
	}
	return 0
}

func getWindowSize(rows *int, cols *int) int {
	var w WinSize
	_, _, err := syscall.Syscall(syscall.SYS_IOCTL,
		os.Stdout.Fd(),
		syscall.TIOCGWINSZ,
		uintptr(unsafe.Pointer(&w)),
	)
	if err != 0 { // type syscall.Errno
		io.WriteString(os.Stdout, "\x1b[999C\x1b[999B")
		return getCursorPosition(rows, cols)
	} else {
		*rows = int(w.Row)
		*cols = int(w.Col)
		return 0
	}
	return -1
}

/*** input ***/

func editorProcessKeypress() {
	c := editorReadKey()
	switch c {
	case ('q' & 0x1f):
		io.WriteString(os.Stdout, "\x1b[2J")
		io.WriteString(os.Stdout, "\x1b[H")
		disableRawMode()
		os.Exit(0)
	}
}

/*** append buffer ***/

type abuf struct {
	buf []byte
}

func (p abuf) String() string {
	return string(p.buf)
}

func (p *abuf) abAppend(s string) {
	p.buf = append(p.buf, []byte(s)...)
}

/*** output ***/

func editorRefreshScreen() {
	var ab abuf
	ab.abAppend("\x1b[25l")
	ab.abAppend("\x1b[H")
	editorDrawRows(&ab)
	ab.abAppend("\x1b[H")
	ab.abAppend("\x1b[25h")
	io.WriteString(os.Stdout, ab.String())
}

func editorDrawRows(ab *abuf) {
	for y := 0; y < E.screenRows-1; y++ {
		if y == E.screenRows/3 {
			w := fmt.Sprintf("Kilo editor -- version %s", KILO_VERSION)
			if len(w) > E.screenCols {
				w = w[0:E.screenCols]
			}
			pad := "~ "
			for padding := (E.screenCols - len(w)) / 2; padding > 0; padding-- {
				ab.abAppend(pad)
				pad = " "
			}
			ab.abAppend(w)
		} else {
			ab.abAppend("~")
		}
		ab.abAppend("\x1b[K")
		if y < E.screenRows-1 {
			ab.abAppend("\r\n")
		}
	}
}

/*** init ***/

func initEditor() {
	if getWindowSize(&E.screenRows, &E.screenCols) == -1 {
		die(fmt.Errorf("couldn't get screen size"))
	}
}

func main() {
	enableRawMode()
	defer disableRawMode()
	initEditor()

	for {
		editorRefreshScreen()
		editorProcessKeypress()
	}
}
<commit_msg>Step 45 - WASD cursor movement<commit_after>package main

import (
	"fmt"
	"io"
	"log"
	"os"
	"syscall"
	"unsafe"
)

/*** defines ***/

const KILO_VERSION = "0.0.1"

/*** data ***/

type Termios struct {
	Iflag  uint32
	Oflag  uint32
	Cflag  uint32
	Lflag  uint32
	Cc     [20]byte
	Ispeed uint32
	Ospeed uint32
}

type editorConfig struct {
	cx          int
	cy          int
	screenRows  int
	screenCols  int
	origTermios *Termios
}

type WinSize struct {
	Row    uint16
	Col    uint16
	Xpixel uint16
	Ypixel uint16
}

var E editorConfig

/*** terminal ***/

func die(err error) {
	disableRawMode()
	io.WriteString(os.Stdout, "\x1b[2J")
	io.WriteString(os.Stdout, "\x1b[H")
	log.Fatal(err)
}

func TcSetAttr(fd uintptr, termios *Termios) error {
	// TCSETS+1 == TCSETSW, because TCSAFLUSH doesn't exist
	if _, _, err := syscall.Syscall(syscall.SYS_IOCTL, fd, uintptr(syscall.TCSETS+1), uintptr(unsafe.Pointer(termios))); err != 0 {
		return err
	}
	return nil
}

func TcGetAttr(fd uintptr) *Termios {
	var termios = &Termios{}
	if _, _, err := syscall.Syscall(syscall.SYS_IOCTL, fd, syscall.TCGETS, uintptr(unsafe.Pointer(termios))); err != 0 {
		log.Fatalf("Problem getting terminal attributes: %s\n", err)
	}
	return termios
}

func enableRawMode() {
	E.origTermios = TcGetAttr(os.Stdin.Fd())
	var raw Termios
	raw = *E.origTermios
	raw.Iflag &^= syscall.BRKINT | syscall.ICRNL | syscall.INPCK | syscall.ISTRIP | syscall.IXON
	raw.Oflag &^= syscall.OPOST
	raw.Cflag |= syscall.CS8
	raw.Lflag &^= syscall.ECHO | syscall.ICANON | syscall.IEXTEN | syscall.ISIG
	raw.Cc[syscall.VMIN+1] = 0
	raw.Cc[syscall.VTIME+1] = 1
	if e := TcSetAttr(os.Stdin.Fd(), &raw); e != nil {
		log.Fatalf("Problem enabling raw mode: %s\n", e)
	}
}

func disableRawMode() {
	if e := TcSetAttr(os.Stdin.Fd(), E.origTermios); e != nil {
		log.Fatalf("Problem disabling raw mode: %s\n", e)
	}
}

func editorReadKey() byte {
	var buffer [1]byte
	var cc int
	var err error
	for cc, err = os.Stdin.Read(buffer[:]); cc != 1; cc, err = os.Stdin.Read(buffer[:]) {
	}
	if err != nil {
		die(err)
	}
	return buffer[0]
}

func getCursorPosition(rows *int, cols *int) int {
	io.WriteString(os.Stdout, "\x1b[6n")
	var buffer [1]byte
	var buf []byte
	var cc int
	for cc, _ = os.Stdin.Read(buffer[:]); cc == 1; cc, _ = os.Stdin.Read(buffer[:]) {
		if buffer[0] == 'R' {
			break
		}
		buf = append(buf, buffer[0])
	}
	if string(buf[0:2]) != "\x1b[" {
		log.Printf("Failed to read rows;cols from tty\n")
		return -1
	}
	if n, e := fmt.Sscanf(string(buf[2:]), "%d;%d", rows, cols); n != 2 || e != nil {
		if e != nil {
			log.Printf("getCursorPosition: fmt.Sscanf() failed: %s\n", e)
		}
		if n != 2 {
			log.Printf("getCursorPosition: got %d items, wanted 2\n", n)
		}
		return -1
	}
	return 0
}

func getWindowSize(rows *int, cols *int) int {
	var w WinSize
	_, _, err := syscall.Syscall(syscall.SYS_IOCTL,
		os.Stdout.Fd(),
		syscall.TIOCGWINSZ,
		uintptr(unsafe.Pointer(&w)),
	)
	if err != 0 { // type syscall.Errno
		io.WriteString(os.Stdout, "\x1b[999C\x1b[999B")
		return getCursorPosition(rows, cols)
	} else {
		*rows = int(w.Row)
		*cols = int(w.Col)
		return 0
	}
	return -1
}

/*** input ***/

func editorMoveCursor(key byte) {
	switch key {
	case 'a':
		E.cx--
	case 'd':
		E.cx++
	case 'w':
		E.cy--
	case 's':
		E.cy++
	}
}

func editorProcessKeypress() {
	c := editorReadKey()
	switch c {
	case ('q' & 0x1f):
		io.WriteString(os.Stdout, "\x1b[2J")
		io.WriteString(os.Stdout, "\x1b[H")
		disableRawMode()
		os.Exit(0)
	case 'w', 'a', 's', 'd':
		editorMoveCursor(c)
	}
}

/*** append buffer ***/

type abuf struct {
	buf []byte
}

func (p abuf) String() string {
	return string(p.buf)
}

func (p *abuf) abAppend(s string) {
	p.buf = append(p.buf, []byte(s)...)
}

/*** output ***/

func editorRefreshScreen() {
	var ab abuf
	ab.abAppend("\x1b[25l")
	ab.abAppend("\x1b[H")
	editorDrawRows(&ab)
	ab.abAppend(fmt.Sprintf("\x1b[%d;%dH", E.cy+1, E.cx+1))
	ab.abAppend("\x1b[25h")
	io.WriteString(os.Stdout, ab.String())
}

func editorDrawRows(ab *abuf) {
	for y := 0; y < E.screenRows-1; y++ {
		if y == E.screenRows/3 {
			w := fmt.Sprintf("Kilo editor -- version %s", KILO_VERSION)
			if len(w) > E.screenCols {
				w = w[0:E.screenCols]
			}
			pad := "~ "
			for padding := (E.screenCols - len(w)) / 2; padding > 0; padding-- {
				ab.abAppend(pad)
				pad = " "
			}
			ab.abAppend(w)
		} else {
			ab.abAppend("~")
		}
		ab.abAppend("\x1b[K")
		if y < E.screenRows-1 {
			ab.abAppend("\r\n")
		}
	}
}

/*** init ***/

func initEditor() {
	if getWindowSize(&E.screenRows, &E.screenCols) == -1 {
		die(fmt.Errorf("couldn't get screen size"))
	}
}

func main() {
	enableRawMode()
	defer disableRawMode()
	initEditor()

	for {
		editorRefreshScreen()
		editorProcessKeypress()
	}
}
<|endoftext|>