<commit_before>// eXtended ETHereum
package xeth

import (
	"bytes"
	"encoding/json"
	"fmt"
	"math/big"

	"github.com/ethereum/go-ethereum/accounts"
	"github.com/ethereum/go-ethereum/core"
	"github.com/ethereum/go-ethereum/core/types"
	"github.com/ethereum/go-ethereum/crypto"
	"github.com/ethereum/go-ethereum/ethutil"
	"github.com/ethereum/go-ethereum/event"
	"github.com/ethereum/go-ethereum/logger"
	"github.com/ethereum/go-ethereum/miner"
	"github.com/ethereum/go-ethereum/p2p"
	"github.com/ethereum/go-ethereum/state"
	"github.com/ethereum/go-ethereum/whisper"
)

var pipelogger = logger.NewLogger("XETH")

// to resolve the import cycle
type Backend interface {
	BlockProcessor() *core.BlockProcessor
	ChainManager() *core.ChainManager
	AccountManager() *accounts.Manager
	TxPool() *core.TxPool
	PeerCount() int
	IsListening() bool
	Peers() []*p2p.Peer
	BlockDb() ethutil.Database
	StateDb() ethutil.Database
	EventMux() *event.TypeMux
	Whisper() *whisper.Whisper
	Miner() *miner.Miner
}

// Frontend should be implemented by users of XEth. Its methods are
// called whenever XEth makes a decision that requires user input.
type Frontend interface {
	// UnlockAccount is called when a transaction needs to be signed
	// but the key corresponding to the transaction's sender is
	// locked.
	//
	// It should unlock the account with the given address and return
	// true if unlocking succeeded.
	UnlockAccount(address []byte) bool

	// This is called for all transactions inititated through
	// Transact. It should prompt the user to confirm the transaction
	// and return true if the transaction was acknowledged.
	//
	// ConfirmTransaction is not used for Call transactions
	// because they cannot change any state.
	ConfirmTransaction(tx *types.Transaction) bool
}

type XEth struct {
	eth            Backend
	blockProcessor *core.BlockProcessor
	chainManager   *core.ChainManager
	accountManager *accounts.Manager
	state          *State
	whisper        *Whisper
	miner          *miner.Miner

	frontend Frontend
}

// dummyFrontend is a non-interactive frontend that allows all
// transactions but cannot not unlock any keys.
type dummyFrontend struct{}

func (dummyFrontend) UnlockAccount([]byte) bool                  { return false }
func (dummyFrontend) ConfirmTransaction(*types.Transaction) bool { return true }

// New creates an XEth that uses the given frontend.
// If a nil Frontend is provided, a default frontend which
// confirms all transactions will be used.
func New(eth Backend, frontend Frontend) *XEth {
	xeth := &XEth{
		eth:            eth,
		blockProcessor: eth.BlockProcessor(),
		chainManager:   eth.ChainManager(),
		accountManager: eth.AccountManager(),
		whisper:        NewWhisper(eth.Whisper()),
		miner:          eth.Miner(),
		frontend:       frontend,
	}
	if frontend == nil {
		xeth.frontend = dummyFrontend{}
	}
	xeth.state = NewState(xeth, xeth.chainManager.TransState())
	return xeth
}

func (self *XEth) Backend() Backend { return self.eth }
func (self *XEth) UseState(statedb *state.StateDB) *XEth {
	xeth := &XEth{
		eth:            self.eth,
		blockProcessor: self.blockProcessor,
		chainManager:   self.chainManager,
		whisper:        self.whisper,
		miner:          self.miner,
	}

	xeth.state = NewState(xeth, statedb)
	return xeth
}
func (self *XEth) State() *State { return self.state }

func (self *XEth) Whisper() *Whisper   { return self.whisper }
func (self *XEth) Miner() *miner.Miner { return self.miner }

func (self *XEth) BlockByHash(strHash string) *Block {
	hash := fromHex(strHash)
	block := self.chainManager.GetBlock(hash)

	return NewBlock(block)
}

func (self *XEth) BlockByNumber(num int32) *Block {
	if num == -1 {
		return NewBlock(self.chainManager.CurrentBlock())
	}

	return NewBlock(self.chainManager.GetBlockByNumber(uint64(num)))
}

func (self *XEth) Block(v interface{}) *Block {
	if n, ok := v.(int32); ok {
		return self.BlockByNumber(n)
	} else if str, ok := v.(string); ok {
		return self.BlockByHash(str)
	} else if f, ok := v.(float64); ok { // Don't ask ...
		return self.BlockByNumber(int32(f))
	}

	return nil
}

func (self *XEth) Accounts() []string {
	// TODO: check err?
	accounts, _ := self.eth.AccountManager().Accounts()
	accountAddresses := make([]string, len(accounts))
	for i, ac := range accounts {
		accountAddresses[i] = toHex(ac.Address)
	}
	return accountAddresses
}

func (self *XEth) PeerCount() int {
	return self.eth.PeerCount()
}

func (self *XEth) IsMining() bool {
	return self.miner.Mining()
}

func (self *XEth) SetMining(shouldmine bool) bool {
	ismining := self.miner.Mining()
	if shouldmine && !ismining {
		self.miner.Start()
	}
	if ismining && !shouldmine {
		self.miner.Stop()
	}
	return self.miner.Mining()
}

func (self *XEth) IsListening() bool {
	return self.eth.IsListening()
}

func (self *XEth) Coinbase() string {
	cb, _ := self.eth.AccountManager().Coinbase()
	return toHex(cb)
}

func (self *XEth) NumberToHuman(balance string) string {
	b := ethutil.Big(balance)

	return ethutil.CurrencyToString(b)
}

func (self *XEth) StorageAt(addr, storageAddr string) string {
	storage := self.State().SafeGet(addr).StorageString(storageAddr)

	return toHex(storage.Bytes())
}

func (self *XEth) BalanceAt(addr string) string {
	return self.State().SafeGet(addr).Balance().String()
}

func (self *XEth) TxCountAt(address string) int {
	return int(self.State().SafeGet(address).Nonce())
}

func (self *XEth) CodeAt(address string) string {
	return toHex(self.State().SafeGet(address).Code())
}

func (self *XEth) IsContract(address string) bool {
	return len(self.State().SafeGet(address).Code()) > 0
}

func (self *XEth) SecretToAddress(key string) string {
	pair, err := crypto.NewKeyPairFromSec(fromHex(key))
	if err != nil {
		return ""
	}

	return toHex(pair.Address())
}

func (self *XEth) Execute(addr, value, gas, price, data string) (string, error) {
	return "", nil
}

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

func (self *XEth) EachStorage(addr string) string {
	var values []KeyVal
	object := self.State().SafeGet(addr)
	it := object.Trie().Iterator()
	for it.Next() {
		values = append(values, KeyVal{toHex(it.Key), toHex(it.Value)})
	}

	valuesJson, err := json.Marshal(values)
	if err != nil {
		return ""
	}

	return string(valuesJson)
}

func (self *XEth) ToAscii(str string) string {
	padded := ethutil.RightPadBytes([]byte(str), 32)

	return "0x" + toHex(padded)
}

func (self *XEth) FromAscii(str string) string {
	if ethutil.IsHex(str) {
		str = str[2:]
	}

	return string(bytes.Trim(fromHex(str), "\x00"))
}

func (self *XEth) FromNumber(str string) string {
	if ethutil.IsHex(str) {
		str = str[2:]
	}

	return ethutil.BigD(fromHex(str)).String()
}

func (self *XEth) PushTx(encodedTx string) (string, error) {
	tx := types.NewTransactionFromBytes(fromHex(encodedTx))
	err := self.eth.TxPool().Add(tx)
	if err != nil {
		return "", err
	}

	if tx.To() == nil {
		addr := core.AddressFromMessage(tx)
		return toHex(addr), nil
	}
	return toHex(tx.Hash()), nil
}

func (self *XEth) Call(fromStr, toStr, valueStr, gasStr, gasPriceStr, dataStr string) (string, error) {
	if len(gasStr) == 0 {
		gasStr = "100000"
	}
	if len(gasPriceStr) == 0 {
		gasPriceStr = "1"
	}

	statedb := self.State().State() //self.chainManager.TransState()
	msg := callmsg{
		from:     statedb.GetOrNewStateObject(fromHex(fromStr)),
		to:       fromHex(toStr),
		gas:      ethutil.Big(gasStr),
		gasPrice: ethutil.Big(gasPriceStr),
		value:    ethutil.Big(valueStr),
		data:     fromHex(dataStr),
	}
	block := self.chainManager.CurrentBlock()
	vmenv := core.NewEnv(statedb, self.chainManager, msg, block)

	res, err := vmenv.Call(msg.from, msg.to, msg.data, msg.gas, msg.gasPrice, msg.value)
	return toHex(res), err
}

func (self *XEth) Transact(fromStr, toStr, valueStr, gasStr, gasPriceStr, codeStr string) (string, error) {
	var (
		from             []byte
		to               []byte
		value            = ethutil.NewValue(valueStr)
		gas              = ethutil.NewValue(gasStr)
		price            = ethutil.NewValue(gasPriceStr)
		data             []byte
		contractCreation bool
	)

	from = fromHex(fromStr)
	data = fromHex(codeStr)
	to = fromHex(toStr)
	if len(to) == 0 {
		contractCreation = true
	}

	var tx *types.Transaction
	if contractCreation {
		tx = types.NewContractCreationTx(value.BigInt(), gas.BigInt(), price.BigInt(), data)
	} else {
		tx = types.NewTransactionMessage(to, value.BigInt(), gas.BigInt(), price.BigInt(), data)
	}

	state := self.chainManager.TransState()
	nonce := state.GetNonce(from)
	tx.SetNonce(nonce)

	if err := self.sign(tx, from, false); err != nil {
		return "", err
	}
	if err := self.eth.TxPool().Add(tx); err != nil {
		return "", err
	}
	state.SetNonce(from, nonce+1)

	if contractCreation {
		addr := core.AddressFromMessage(tx)
		pipelogger.Infof("Contract addr %x\n", addr)
	}

	if types.IsContractAddr(to) {
		return toHex(core.AddressFromMessage(tx)), nil
	}
	return toHex(tx.Hash()), nil
}

func (self *XEth) sign(tx *types.Transaction, from []byte, didUnlock bool) error {
	sig, err := self.accountManager.Sign(accounts.Account{Address: from}, tx.Hash())
	if err == accounts.ErrLocked {
		if didUnlock {
			return fmt.Errorf("sender account still locked after successful unlock")
		}
		if !self.frontend.UnlockAccount(from) {
			return fmt.Errorf("could not unlock sender account")
		}
		// retry signing, the account should now be unlocked.
		self.sign(tx, from, true)
	} else if err != nil {
		return err
	}
	tx.SetSignatureValues(sig)
	return nil
}

// callmsg is the message type used for call transations.
type callmsg struct {
	from          *state.StateObject
	to            []byte
	gas, gasPrice *big.Int
	value         *big.Int
	data          []byte
}

// accessor boilerplate to implement core.Message
func (m callmsg) From() []byte       { return m.from.Address() }
func (m callmsg) Nonce() uint64      { return m.from.Nonce() }
func (m callmsg) To() []byte         { return m.to }
func (m callmsg) GasPrice() *big.Int { return m.gasPrice }
func (m callmsg) Gas() *big.Int      { return m.gas }
func (m callmsg) Value() *big.Int    { return m.value }
func (m callmsg) Data() []byte       { return m.data }
<commit_msg>xeth: fix signing transaction after unlock<commit_after>// eXtended ETHereum
package xeth

import (
	"bytes"
	"encoding/json"
	"fmt"
	"math/big"

	"github.com/ethereum/go-ethereum/accounts"
	"github.com/ethereum/go-ethereum/core"
	"github.com/ethereum/go-ethereum/core/types"
	"github.com/ethereum/go-ethereum/crypto"
	"github.com/ethereum/go-ethereum/ethutil"
	"github.com/ethereum/go-ethereum/event"
	"github.com/ethereum/go-ethereum/logger"
	"github.com/ethereum/go-ethereum/miner"
	"github.com/ethereum/go-ethereum/p2p"
	"github.com/ethereum/go-ethereum/state"
	"github.com/ethereum/go-ethereum/whisper"
)

var pipelogger = logger.NewLogger("XETH")

// to resolve the import cycle
type Backend interface {
	BlockProcessor() *core.BlockProcessor
	ChainManager() *core.ChainManager
	AccountManager() *accounts.Manager
	TxPool() *core.TxPool
	PeerCount() int
	IsListening() bool
	Peers() []*p2p.Peer
	BlockDb() ethutil.Database
	StateDb() ethutil.Database
	EventMux() *event.TypeMux
	Whisper() *whisper.Whisper
	Miner() *miner.Miner
}

// Frontend should be implemented by users of XEth. Its methods are
// called whenever XEth makes a decision that requires user input.
type Frontend interface {
	// UnlockAccount is called when a transaction needs to be signed
	// but the key corresponding to the transaction's sender is
	// locked.
	//
	// It should unlock the account with the given address and return
	// true if unlocking succeeded.
	UnlockAccount(address []byte) bool

	// This is called for all transactions inititated through
	// Transact. It should prompt the user to confirm the transaction
	// and return true if the transaction was acknowledged.
	//
	// ConfirmTransaction is not used for Call transactions
	// because they cannot change any state.
	ConfirmTransaction(tx *types.Transaction) bool
}

type XEth struct {
	eth            Backend
	blockProcessor *core.BlockProcessor
	chainManager   *core.ChainManager
	accountManager *accounts.Manager
	state          *State
	whisper        *Whisper
	miner          *miner.Miner

	frontend Frontend
}

// dummyFrontend is a non-interactive frontend that allows all
// transactions but cannot not unlock any keys.
type dummyFrontend struct{}

func (dummyFrontend) UnlockAccount([]byte) bool                  { return false }
func (dummyFrontend) ConfirmTransaction(*types.Transaction) bool { return true }

// New creates an XEth that uses the given frontend.
// If a nil Frontend is provided, a default frontend which
// confirms all transactions will be used.
func New(eth Backend, frontend Frontend) *XEth {
	xeth := &XEth{
		eth:            eth,
		blockProcessor: eth.BlockProcessor(),
		chainManager:   eth.ChainManager(),
		accountManager: eth.AccountManager(),
		whisper:        NewWhisper(eth.Whisper()),
		miner:          eth.Miner(),
		frontend:       frontend,
	}
	if frontend == nil {
		xeth.frontend = dummyFrontend{}
	}
	xeth.state = NewState(xeth, xeth.chainManager.TransState())
	return xeth
}

func (self *XEth) Backend() Backend { return self.eth }
func (self *XEth) UseState(statedb *state.StateDB) *XEth {
	xeth := &XEth{
		eth:            self.eth,
		blockProcessor: self.blockProcessor,
		chainManager:   self.chainManager,
		whisper:        self.whisper,
		miner:          self.miner,
	}

	xeth.state = NewState(xeth, statedb)
	return xeth
}
func (self *XEth) State() *State { return self.state }

func (self *XEth) Whisper() *Whisper   { return self.whisper }
func (self *XEth) Miner() *miner.Miner { return self.miner }

func (self *XEth) BlockByHash(strHash string) *Block {
	hash := fromHex(strHash)
	block := self.chainManager.GetBlock(hash)

	return NewBlock(block)
}

func (self *XEth) BlockByNumber(num int32) *Block {
	if num == -1 {
		return NewBlock(self.chainManager.CurrentBlock())
	}

	return NewBlock(self.chainManager.GetBlockByNumber(uint64(num)))
}

func (self *XEth) Block(v interface{}) *Block {
	if n, ok := v.(int32); ok {
		return self.BlockByNumber(n)
	} else if str, ok := v.(string); ok {
		return self.BlockByHash(str)
	} else if f, ok := v.(float64); ok { // Don't ask ...
		return self.BlockByNumber(int32(f))
	}

	return nil
}

func (self *XEth) Accounts() []string {
	// TODO: check err?
	accounts, _ := self.eth.AccountManager().Accounts()
	accountAddresses := make([]string, len(accounts))
	for i, ac := range accounts {
		accountAddresses[i] = toHex(ac.Address)
	}
	return accountAddresses
}

func (self *XEth) PeerCount() int {
	return self.eth.PeerCount()
}

func (self *XEth) IsMining() bool {
	return self.miner.Mining()
}

func (self *XEth) SetMining(shouldmine bool) bool {
	ismining := self.miner.Mining()
	if shouldmine && !ismining {
		self.miner.Start()
	}
	if ismining && !shouldmine {
		self.miner.Stop()
	}
	return self.miner.Mining()
}

func (self *XEth) IsListening() bool {
	return self.eth.IsListening()
}

func (self *XEth) Coinbase() string {
	cb, _ := self.eth.AccountManager().Coinbase()
	return toHex(cb)
}

func (self *XEth) NumberToHuman(balance string) string {
	b := ethutil.Big(balance)

	return ethutil.CurrencyToString(b)
}

func (self *XEth) StorageAt(addr, storageAddr string) string {
	storage := self.State().SafeGet(addr).StorageString(storageAddr)

	return toHex(storage.Bytes())
}

func (self *XEth) BalanceAt(addr string) string {
	return self.State().SafeGet(addr).Balance().String()
}

func (self *XEth) TxCountAt(address string) int {
	return int(self.State().SafeGet(address).Nonce())
}

func (self *XEth) CodeAt(address string) string {
	return toHex(self.State().SafeGet(address).Code())
}

func (self *XEth) IsContract(address string) bool {
	return len(self.State().SafeGet(address).Code()) > 0
}

func (self *XEth) SecretToAddress(key string) string {
	pair, err := crypto.NewKeyPairFromSec(fromHex(key))
	if err != nil {
		return ""
	}

	return toHex(pair.Address())
}

func (self *XEth) Execute(addr, value, gas, price, data string) (string, error) {
	return "", nil
}

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

func (self *XEth) EachStorage(addr string) string {
	var values []KeyVal
	object := self.State().SafeGet(addr)
	it := object.Trie().Iterator()
	for it.Next() {
		values = append(values, KeyVal{toHex(it.Key), toHex(it.Value)})
	}

	valuesJson, err := json.Marshal(values)
	if err != nil {
		return ""
	}

	return string(valuesJson)
}

func (self *XEth) ToAscii(str string) string {
	padded := ethutil.RightPadBytes([]byte(str), 32)

	return "0x" + toHex(padded)
}

func (self *XEth) FromAscii(str string) string {
	if ethutil.IsHex(str) {
		str = str[2:]
	}

	return string(bytes.Trim(fromHex(str), "\x00"))
}

func (self *XEth) FromNumber(str string) string {
	if ethutil.IsHex(str) {
		str = str[2:]
	}

	return ethutil.BigD(fromHex(str)).String()
}

func (self *XEth) PushTx(encodedTx string) (string, error) {
	tx := types.NewTransactionFromBytes(fromHex(encodedTx))
	err := self.eth.TxPool().Add(tx)
	if err != nil {
		return "", err
	}

	if tx.To() == nil {
		addr := core.AddressFromMessage(tx)
		return toHex(addr), nil
	}
	return toHex(tx.Hash()), nil
}

func (self *XEth) Call(fromStr, toStr, valueStr, gasStr, gasPriceStr, dataStr string) (string, error) {
	if len(gasStr) == 0 {
		gasStr = "100000"
	}
	if len(gasPriceStr) == 0 {
		gasPriceStr = "1"
	}

	statedb := self.State().State() //self.chainManager.TransState()
	msg := callmsg{
		from:     statedb.GetOrNewStateObject(fromHex(fromStr)),
		to:       fromHex(toStr),
		gas:      ethutil.Big(gasStr),
		gasPrice: ethutil.Big(gasPriceStr),
		value:    ethutil.Big(valueStr),
		data:     fromHex(dataStr),
	}
	block := self.chainManager.CurrentBlock()
	vmenv := core.NewEnv(statedb, self.chainManager, msg, block)

	res, err := vmenv.Call(msg.from, msg.to, msg.data, msg.gas, msg.gasPrice, msg.value)
	return toHex(res), err
}

func (self *XEth) Transact(fromStr, toStr, valueStr, gasStr, gasPriceStr, codeStr string) (string, error) {
	var (
		from             []byte
		to               []byte
		value            = ethutil.NewValue(valueStr)
		gas              = ethutil.NewValue(gasStr)
		price            = ethutil.NewValue(gasPriceStr)
		data             []byte
		contractCreation bool
	)

	from = fromHex(fromStr)
	data = fromHex(codeStr)
	to = fromHex(toStr)
	if len(to) == 0 {
		contractCreation = true
	}

	var tx *types.Transaction
	if contractCreation {
		tx = types.NewContractCreationTx(value.BigInt(), gas.BigInt(), price.BigInt(), data)
	} else {
		tx = types.NewTransactionMessage(to, value.BigInt(), gas.BigInt(), price.BigInt(), data)
	}

	state := self.chainManager.TransState()
	nonce := state.GetNonce(from)
	tx.SetNonce(nonce)

	if err := self.sign(tx, from, false); err != nil {
		return "", err
	}
	if err := self.eth.TxPool().Add(tx); err != nil {
		return "", err
	}
	state.SetNonce(from, nonce+1)

	if contractCreation {
		addr := core.AddressFromMessage(tx)
		pipelogger.Infof("Contract addr %x\n", addr)
	}

	if types.IsContractAddr(to) {
		return toHex(core.AddressFromMessage(tx)), nil
	}
	return toHex(tx.Hash()), nil
}

func (self *XEth) sign(tx *types.Transaction, from []byte, didUnlock bool) error {
	sig, err := self.accountManager.Sign(accounts.Account{Address: from}, tx.Hash())
	if err == accounts.ErrLocked {
		if didUnlock {
			return fmt.Errorf("sender account still locked after successful unlock")
		}
		if !self.frontend.UnlockAccount(from) {
			return fmt.Errorf("could not unlock sender account")
		}
		// retry signing, the account should now be unlocked.
		return self.sign(tx, from, true)
	} else if err != nil {
		return err
	}
	tx.SetSignatureValues(sig)
	return nil
}

// callmsg is the message type used for call transations.
type callmsg struct {
	from          *state.StateObject
	to            []byte
	gas, gasPrice *big.Int
	value         *big.Int
	data          []byte
}

// accessor boilerplate to implement core.Message
func (m callmsg) From() []byte       { return m.from.Address() }
func (m callmsg) Nonce() uint64      { return m.from.Nonce() }
func (m callmsg) To() []byte         { return m.to }
func (m callmsg) GasPrice() *big.Int { return m.gasPrice }
func (m callmsg) Gas() *big.Int      { return m.gas }
func (m callmsg) Value() *big.Int    { return m.value }
func (m callmsg) Data() []byte       { return m.data }
<|endoftext|>


<commit_before>package ydls

import (
	"context"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"sort"
	"sync"

	"github.com/wader/ydls/ffmpeg"
	"github.com/wader/ydls/id3v2"
	"github.com/wader/ydls/rereadcloser"
	"github.com/wader/ydls/writelogger"
	"github.com/wader/ydls/youtubedl"
)

func firstNonEmpty(sl ...string) string {
	for _, s := range sl {
		if s != "" {
			return s
		}
	}
	return ""
}

func logOrDiscard(l *log.Logger) *log.Logger {
	if l != nil {
		return l
	}

	return log.New(ioutil.Discard, "", 0)
}

func id3v2FramesFromYoutueDLInfo(i *youtubedl.Info) []id3v2.Frame {
	frames := []id3v2.Frame{
		&id3v2.TextFrame{ID: "TPE1", Text: firstNonEmpty(i.Artist, i.Creator, i.Uploader)},
		&id3v2.TextFrame{ID: "TIT2", Text: i.Title},
		&id3v2.COMMFrame{Language: "XXX", Description: "", Text: i.Description},
	}
	if i.Duration > 0 {
		frames = append(frames, &id3v2.TextFrame{
			ID:   "TLEN",
			Text: fmt.Sprintf("%d", uint32(i.Duration*1000)),
		})
	}
	if len(i.ThumbnailBytes) > 0 {
		frames = append(frames, &id3v2.APICFrame{
			MIMEType:    http.DetectContentType(i.ThumbnailBytes),
			PictureType: id3v2.PictureTypeOther,
			Description: "",
			Data:        i.ThumbnailBytes,
		})
	}

	return frames
}

func findFormat(formats []*youtubedl.Format, protocol string, aCodecs *prioStringSet, vCodecs *prioStringSet) *youtubedl.Format {
	var matched []*youtubedl.Format

	for _, f := range formats {
		if protocol != "*" && f.Protocol != protocol {
			continue
		}
		if !(aCodecs == nil || (f.NormACodec == "" && aCodecs.empty()) || aCodecs.member(f.NormACodec)) {
			continue
		}
		if !(vCodecs == nil || (f.NormVCodec == "" && vCodecs.empty()) || vCodecs.member(f.NormVCodec)) {
			continue
		}

		matched = append(matched, f)
	}

	sort.Sort(youtubedl.FormatByNormBR(matched))

	if len(matched) > 0 {
		return matched[0]
	}

	return nil
}

func findBestFormats(ydlFormats []*youtubedl.Format, format *Format) (aFormat *youtubedl.Format, vFormat *youtubedl.Format) {
	type neededFormat struct {
		aCodecs    *prioStringSet
		vCodecs    *prioStringSet
		aYDLFormat **youtubedl.Format
		vYDLFormat **youtubedl.Format
	}

	// TODO: messy, needs rewrite

	var neededFormats []*neededFormat

	// match exactly, both audio/video codecs found or not found
	neededFormats = append(neededFormats, &neededFormat{
		format.ACodecs.PrioStringSet(),
		format.VCodecs.PrioStringSet(),
		&aFormat, &vFormat,
	})

	if !format.ACodecs.empty() {
		// matching audio codec with any video codec
		neededFormats = append(neededFormats, &neededFormat{format.ACodecs.PrioStringSet(), nil, &aFormat, nil})
		// match any audio codec and no video
		neededFormats = append(neededFormats, &neededFormat{nil, &prioStringSet{}, &aFormat, nil})
		// match any audio and video codec
		neededFormats = append(neededFormats, &neededFormat{nil, nil, &aFormat, nil})
	}
	if !format.VCodecs.empty() {
		// same logic as above
		neededFormats = append(neededFormats, &neededFormat{nil, format.VCodecs.PrioStringSet(), nil, &vFormat})
		neededFormats = append(neededFormats, &neededFormat{&prioStringSet{}, nil, nil, &vFormat})
		neededFormats = append(neededFormats, &neededFormat{nil, nil, nil, &vFormat})
	}

	// TODO: if only audio => stream with lowest video br?

	for _, proto := range []string{"https", "http", "*"} {
		for _, f := range neededFormats {
			m := findFormat(ydlFormats, proto, f.aCodecs, f.vCodecs)

			if m == nil {
				continue
			}

			if f.aYDLFormat != nil && *f.aYDLFormat == nil && m.NormACodec != "" {
				*f.aYDLFormat = m
			}
			if f.vYDLFormat != nil && *f.vYDLFormat == nil && m.NormVCodec != "" {
				*f.vYDLFormat = m
			}

			if (format.ACodecs.empty() || aFormat != nil) &&
				(format.VCodecs.empty() || vFormat != nil) {
				break
			}
		}
	}

	return aFormat, vFormat
}

func downloadAndProbeFormat(ctx context.Context, ydl *youtubedl.Info, filter string, debugLog *log.Logger) (r io.ReadCloser, pi *ffmpeg.ProbeInfo, err error) {
	log := logOrDiscard(debugLog)

	ydlStderr := writelogger.New(log, fmt.Sprintf("ydl-dl %s stderr> ", filter))
	r, err = ydl.Download(ctx, filter, ydlStderr)
	if err != nil {
		return nil, nil, err
	}

	rr := rereadcloser.New(r)
	ffprobeStderr := writelogger.New(log, fmt.Sprintf("ffprobe %s stderr> ", filter))
	const maxProbeByteSize = 10 * 1024 * 1024
	pi, err = ffmpeg.Probe(ctx, io.LimitReader(rr, maxProbeByteSize), log, ffprobeStderr)
	if err != nil {
		return nil, nil, err
	}
	// restart and replay buffer data used when probing
	rr.Restarted = true

	return rr, pi, nil
}

// YDLS youtubedl downloader with some extras
type YDLS struct {
	Formats *Formats
}

// NewFromFile new YDLs using formats file
func NewFromFile(formatsPath string) (*YDLS, error) {
	formatsFile, err := os.Open(formatsPath)
	if err != nil {
		return nil, err
	}
	defer formatsFile.Close()
	formats, err := parseFormats(formatsFile)
	if err != nil {
		return nil, err
	}

	return &YDLS{Formats: formats}, nil
}

// DownloadResult download result
type DownloadResult struct {
	Media    io.ReadCloser
	Filename string
	MIMEType string
}

func chooseFormatCodec(formats prioFormatCodecSet, probedCodec string) *FormatCodec {
	codecFormat := formats.findByCodecName(probedCodec)
	if codecFormat != nil {
		copyCodecFormat := *codecFormat
		copyCodecFormat.Codec = "copy"
		return &copyCodecFormat
	}

	return formats.first()
}

func fancyYDLFormatName(ydlFormat *youtubedl.Format) string {
	if ydlFormat == nil {
		return "n/a"
	}
	return ydlFormat.String()
}

// Download downloads media from URL using context and makes sure output is in specified format
func (ydls *YDLS) Download(ctx context.Context, url string, formatName string, debugLog *log.Logger) (*DownloadResult, error) {
	var closeOnDone []io.Closer
	closeOnDoneFn := func() {
		for _, c := range closeOnDone {
			c.Close()
		}
	}
	deferCloseFn := closeOnDoneFn
	defer func() {
		// will be nil if cmd starts and goroutine takes care of closing instead
		if deferCloseFn != nil {
			deferCloseFn()
		}
	}()

	log := logOrDiscard(debugLog)

	log.Printf("URL: %s", url)
	log.Printf("Output format: %s", formatName)

	var ydlStdout io.Writer
	ydlStdout = writelogger.New(log, "ydl-new stdout> ")
	ydl, err := youtubedl.NewFromURL(ctx, url, ydlStdout)
	if err != nil {
		log.Printf("Failed to download: %s", err)
		return nil, err
	}

	log.Printf("Title: %s", ydl.Title)
	log.Printf("Available youtubedl formats:")
	for _, f := range ydl.Formats {
		log.Printf("  %s", f)
	}

	dr := &DownloadResult{}

	if formatName == "" {
		var probeInfo *ffmpeg.ProbeInfo
		dr.Media, probeInfo, err = downloadAndProbeFormat(ctx, ydl, "best[protocol=https]/best[protocol=http]/best", log)
		if err != nil {
			return nil, err
		}

		log.Printf("Probed format %s", probeInfo)

		// see if we know about the probed format, otherwise fallback to "raw"
		outFormat := ydls.Formats.Find(probeInfo.FormatName(), probeInfo.ACodec(), probeInfo.VCodec())
		dr.MIMEType = "application/octet-stream"
		if outFormat != nil {
			dr.MIMEType = outFormat.MIMEType
		}
		dr.Filename = ydl.Title + ".raw"

		return dr, nil
	}

	outFormat := ydls.Formats.FindByName(formatName)
	if outFormat == nil {
		return nil, fmt.Errorf("could not find format")
	}

	dr.MIMEType = outFormat.MIMEType
	dr.Filename = ydl.Title + "." + outFormat.Ext

	aYDLFormat, vYDLFormat := findBestFormats(ydl.Formats, outFormat)

	var aProbeInfo *ffmpeg.ProbeInfo
	var aReader io.ReadCloser
	var aErr error
	var vProbeInfo *ffmpeg.ProbeInfo
	var vReader io.ReadCloser
	var vErr error

	if aYDLFormat != nil && vYDLFormat != nil {
		if aYDLFormat != vYDLFormat {
			// audio and video in different formats, download simultaneously
			var probeWG sync.WaitGroup
			probeWG.Add(2)
			go func() {
				defer probeWG.Done()
				aReader, aProbeInfo, aErr = downloadAndProbeFormat(ctx, ydl, aYDLFormat.FormatID, log)
			}()
			go func() {
				defer probeWG.Done()
				vReader, vProbeInfo, vErr = downloadAndProbeFormat(ctx, ydl, vYDLFormat.FormatID, log)
			}()
			probeWG.Wait()
			if aReader != nil {
				closeOnDone = append(closeOnDone, aReader)
			}
			if vReader != nil {
				closeOnDone = append(closeOnDone, vReader)
			}
		} else {
			// audio and video in same format
			aReader, aProbeInfo, aErr = downloadAndProbeFormat(ctx, ydl, aYDLFormat.FormatID, log)
			vReader, vProbeInfo, vErr = aReader, aProbeInfo, aErr
			if aReader != nil {
				closeOnDone = append(closeOnDone, aReader)
			}
		}
	} else if aYDLFormat != nil && vYDLFormat == nil {
		// only audio format
		aReader, aProbeInfo, aErr = downloadAndProbeFormat(ctx, ydl, aYDLFormat.FormatID, log)
		if aReader != nil {
			closeOnDone = append(closeOnDone, aReader)
		}
	} else {
		// don't know, download and probe
		aReader, aProbeInfo, aErr = downloadAndProbeFormat(ctx, ydl, "best", log)
		vReader, vProbeInfo, vErr = aReader, aProbeInfo, aErr
		if aReader != nil {
			closeOnDone = append(closeOnDone, aReader)
		}
	}
	if aErr != nil || vErr != nil {
		return nil, fmt.Errorf("failed to probe")
	}

	log.Printf("Stream mapping:")

	var streamMaps []ffmpeg.StreamMap
	ffmpegFormatFlags := outFormat.FormatFlags

	if len(outFormat.ACodecs) > 0 && aProbeInfo != nil && aProbeInfo.ACodec() != "" {
		codecFormat := chooseFormatCodec(outFormat.ACodecs, aProbeInfo.ACodec())
		streamMaps = append(streamMaps, ffmpeg.StreamMap{
			Reader:     aReader,
			Specifier:  "a:0",
			Codec:      "acodec:" + codecFormat.Codec,
			CodecFlags: codecFormat.CodecFlags,
		})
		ffmpegFormatFlags = append(ffmpegFormatFlags, codecFormat.FormatFlags...)

		log.Printf("  audio %s probed:%s -> %s",
			fancyYDLFormatName(aYDLFormat),
			aProbeInfo,
			codecFormat.Codec,
		)
	}
	if len(outFormat.VCodecs) > 0 && vProbeInfo != nil && vProbeInfo.VCodec() != "" {
		codecFormat := chooseFormatCodec(outFormat.VCodecs, vProbeInfo.VCodec())
		streamMaps = append(streamMaps, ffmpeg.StreamMap{
			Reader:     vReader,
			Specifier:  "v:0",
			Codec:      "vcodec:" + codecFormat.Codec,
			CodecFlags: codecFormat.CodecFlags,
		})
		ffmpegFormatFlags = append(ffmpegFormatFlags, codecFormat.FormatFlags...)

		log.Printf("  video %s probed:%s -> %s",
			fancyYDLFormatName(vYDLFormat),
			vProbeInfo,
			codecFormat.Codec,
		)
	}

	var ffmpegStderr io.Writer
	ffmpegStderr = writelogger.New(log, "ffmpeg stderr> ")
	ffmpegR, ffmpegW := io.Pipe()
	closeOnDone = append(closeOnDone, ffmpegR)

	ffmpegP := &ffmpeg.FFmpeg{
		StreamMaps: streamMaps,
		Format:     ffmpeg.Format{Name: outFormat.Formats.first(), Flags: ffmpegFormatFlags},
		DebugLog:   log,
		Stdout:     ffmpegW,
		Stderr:     ffmpegStderr,
	}

	if err := ffmpegP.Start(ctx); err != nil {
		return nil, err
	}

	// goroutine will take care of closing
	deferCloseFn = nil

	var w io.WriteCloser
	dr.Media, w = io.Pipe()
	closeOnDone = append(closeOnDone, w)

	go func() {
		if outFormat.Prepend == "id3v2" {
			id3v2.Write(w, id3v2FramesFromYoutueDLInfo(ydl))
		}
		io.Copy(w, ffmpegR)
		closeOnDoneFn()
		ffmpegP.Wait()
	}()

	return dr, nil
}
<commit_msg>For raw use filename ext if format is known<commit_after>package ydls

import (
	"context"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"sort"
	"sync"

	"github.com/wader/ydls/ffmpeg"
	"github.com/wader/ydls/id3v2"
	"github.com/wader/ydls/rereadcloser"
	"github.com/wader/ydls/writelogger"
	"github.com/wader/ydls/youtubedl"
)

func firstNonEmpty(sl ...string) string {
	for _, s := range sl {
		if s != "" {
			return s
		}
	}
	return ""
}

func logOrDiscard(l *log.Logger) *log.Logger {
	if l != nil {
		return l
	}

	return log.New(ioutil.Discard, "", 0)
}

func id3v2FramesFromYoutueDLInfo(i *youtubedl.Info) []id3v2.Frame {
	frames := []id3v2.Frame{
		&id3v2.TextFrame{ID: "TPE1", Text: firstNonEmpty(i.Artist, i.Creator, i.Uploader)},
		&id3v2.TextFrame{ID: "TIT2", Text: i.Title},
		&id3v2.COMMFrame{Language: "XXX", Description: "", Text: i.Description},
	}
	if i.Duration > 0 {
		frames = append(frames, &id3v2.TextFrame{
			ID:   "TLEN",
			Text: fmt.Sprintf("%d", uint32(i.Duration*1000)),
		})
	}
	if len(i.ThumbnailBytes) > 0 {
		frames = append(frames, &id3v2.APICFrame{
			MIMEType:    http.DetectContentType(i.ThumbnailBytes),
			PictureType: id3v2.PictureTypeOther,
			Description: "",
			Data:        i.ThumbnailBytes,
		})
	}

	return frames
}

func findFormat(formats []*youtubedl.Format, protocol string, aCodecs *prioStringSet, vCodecs *prioStringSet) *youtubedl.Format {
	var matched []*youtubedl.Format

	for _, f := range formats {
		if protocol != "*" && f.Protocol != protocol {
			continue
		}
		if !(aCodecs == nil || (f.NormACodec == "" && aCodecs.empty()) || aCodecs.member(f.NormACodec)) {
			continue
		}
		if !(vCodecs == nil || (f.NormVCodec == "" && vCodecs.empty()) || vCodecs.member(f.NormVCodec)) {
			continue
		}

		matched = append(matched, f)
	}

	sort.Sort(youtubedl.FormatByNormBR(matched))

	if len(matched) > 0 {
		return matched[0]
	}

	return nil
}

func findBestFormats(ydlFormats []*youtubedl.Format, format *Format) (aFormat *youtubedl.Format, vFormat *youtubedl.Format) {
	type neededFormat struct {
		aCodecs    *prioStringSet
		vCodecs    *prioStringSet
		aYDLFormat **youtubedl.Format
		vYDLFormat **youtubedl.Format
	}

	// TODO: messy, needs rewrite

	var neededFormats []*neededFormat

	// match exactly, both audio/video codecs found or not found
	neededFormats = append(neededFormats, &neededFormat{
		format.ACodecs.PrioStringSet(),
		format.VCodecs.PrioStringSet(),
		&aFormat, &vFormat,
	})

	if !format.ACodecs.empty() {
		// matching audio codec with any video codec
		neededFormats = append(neededFormats, &neededFormat{format.ACodecs.PrioStringSet(), nil, &aFormat, nil})
		// match any audio codec and no video
		neededFormats = append(neededFormats, &neededFormat{nil, &prioStringSet{}, &aFormat, nil})
		// match any audio and video codec
		neededFormats = append(neededFormats, &neededFormat{nil, nil, &aFormat, nil})
	}
	if !format.VCodecs.empty() {
		// same logic as above
		neededFormats = append(neededFormats, &neededFormat{nil, format.VCodecs.PrioStringSet(), nil, &vFormat})
		neededFormats = append(neededFormats, &neededFormat{&prioStringSet{}, nil, nil, &vFormat})
		neededFormats = append(neededFormats, &neededFormat{nil, nil, nil, &vFormat})
	}

	// TODO: if only audio => stream with lowest video br?

	for _, proto := range []string{"https", "http", "*"} {
		for _, f := range neededFormats {
			m := findFormat(ydlFormats, proto, f.aCodecs, f.vCodecs)

			if m == nil {
				continue
			}

			if f.aYDLFormat != nil && *f.aYDLFormat == nil && m.NormACodec != "" {
				*f.aYDLFormat = m
			}
			if f.vYDLFormat != nil && *f.vYDLFormat == nil && m.NormVCodec != "" {
				*f.vYDLFormat = m
			}

			if (format.ACodecs.empty() || aFormat != nil) &&
				(format.VCodecs.empty() || vFormat != nil) {
				break
			}
		}
	}

	return aFormat, vFormat
}

func downloadAndProbeFormat(ctx context.Context, ydl *youtubedl.Info, filter string, debugLog *log.Logger) (r io.ReadCloser, pi *ffmpeg.ProbeInfo, err error) {
	log := logOrDiscard(debugLog)

	ydlStderr := writelogger.New(log, fmt.Sprintf("ydl-dl %s stderr> ", filter))
	r, err = ydl.Download(ctx, filter, ydlStderr)
	if err != nil {
		return nil, nil, err
	}

	rr := rereadcloser.New(r)
	ffprobeStderr := writelogger.New(log, fmt.Sprintf("ffprobe %s stderr> ", filter))
	const maxProbeByteSize = 10 * 1024 * 1024
	pi, err = ffmpeg.Probe(ctx, io.LimitReader(rr, maxProbeByteSize), log, ffprobeStderr)
	if err != nil {
		return nil, nil, err
	}
	// restart and replay buffer data used when probing
	rr.Restarted = true

	return rr, pi, nil
}

// YDLS youtubedl downloader with some extras
type YDLS struct {
	Formats *Formats
}

// NewFromFile new YDLs using formats file
func NewFromFile(formatsPath string) (*YDLS, error) {
	formatsFile, err := os.Open(formatsPath)
	if err != nil {
		return nil, err
	}
	defer formatsFile.Close()
	formats, err := parseFormats(formatsFile)
	if err != nil {
		return nil, err
	}

	return &YDLS{Formats: formats}, nil
}

// DownloadResult download result
type DownloadResult struct {
	Media    io.ReadCloser
	Filename string
	MIMEType string
}

func chooseFormatCodec(formats prioFormatCodecSet, probedCodec string) *FormatCodec {
	codecFormat := formats.findByCodecName(probedCodec)
	if codecFormat != nil {
		copyCodecFormat := *codecFormat
		copyCodecFormat.Codec = "copy"
		return &copyCodecFormat
	}

	return formats.first()
}

func fancyYDLFormatName(ydlFormat *youtubedl.Format) string {
	if ydlFormat == nil {
		return "n/a"
	}
	return ydlFormat.String()
}

// Download downloads media from URL using context and makes sure output is in specified format
func (ydls *YDLS) Download(ctx context.Context, url string, formatName string, debugLog *log.Logger) (*DownloadResult, error) {
	var closeOnDone []io.Closer
	closeOnDoneFn := func() {
		for _, c := range closeOnDone {
			c.Close()
		}
	}
	deferCloseFn := closeOnDoneFn
	defer func() {
		// will be nil if cmd starts and goroutine takes care of closing instead
		if deferCloseFn != nil {
			deferCloseFn()
		}
	}()

	log := logOrDiscard(debugLog)

	log.Printf("URL: %s", url)
	log.Printf("Output format: %s", formatName)

	var ydlStdout io.Writer
	ydlStdout = writelogger.New(log, "ydl-new stdout> ")
	ydl, err := youtubedl.NewFromURL(ctx, url, ydlStdout)
	if err != nil {
		log.Printf("Failed to download: %s", err)
		return nil, err
	}

	log.Printf("Title: %s", ydl.Title)
	log.Printf("Available youtubedl formats:")
	for _, f := range ydl.Formats {
		log.Printf("  %s", f)
	}

	dr := &DownloadResult{}

	if formatName == "" {
		var probeInfo *ffmpeg.ProbeInfo
		dr.Media, probeInfo, err = downloadAndProbeFormat(ctx, ydl, "best[protocol=https]/best[protocol=http]/best", log)
		if err != nil {
			return nil, err
		}

		log.Printf("Probed format %s", probeInfo)

		// see if we know about the probed format, otherwise fallback to "raw"
		outFormat := ydls.Formats.Find(probeInfo.FormatName(), probeInfo.ACodec(), probeInfo.VCodec())
		if outFormat != nil {
			dr.MIMEType = outFormat.MIMEType
			dr.Filename = ydl.Title + "." + outFormat.Ext
		} else {
			dr.MIMEType = "application/octet-stream"
			dr.Filename = ydl.Title + ".raw"
		}

		return dr, nil
	}

	outFormat := ydls.Formats.FindByName(formatName)
	if outFormat == nil {
		return nil, fmt.Errorf("could not find format")
	}

	dr.MIMEType = outFormat.MIMEType
	dr.Filename = ydl.Title + "." + outFormat.Ext

	aYDLFormat, vYDLFormat := findBestFormats(ydl.Formats, outFormat)

	var aProbeInfo *ffmpeg.ProbeInfo
	var aReader io.ReadCloser
	var aErr error
	var vProbeInfo *ffmpeg.ProbeInfo
	var vReader io.ReadCloser
	var vErr error

	if aYDLFormat != nil && vYDLFormat != nil {
		if aYDLFormat != vYDLFormat {
			// audio and video in different formats, download simultaneously
			var probeWG sync.WaitGroup
			probeWG.Add(2)
			go func() {
				defer probeWG.Done()
				aReader, aProbeInfo, aErr = downloadAndProbeFormat(ctx, ydl, aYDLFormat.FormatID, log)
			}()
			go func() {
				defer probeWG.Done()
				vReader, vProbeInfo, vErr = downloadAndProbeFormat(ctx, ydl, vYDLFormat.FormatID, log)
			}()
			probeWG.Wait()
			if aReader != nil {
				closeOnDone = append(closeOnDone, aReader)
			}
			if vReader != nil {
				closeOnDone = append(closeOnDone, vReader)
			}
		} else {
			// audio and video in same format
			aReader, aProbeInfo, aErr = downloadAndProbeFormat(ctx, ydl, aYDLFormat.FormatID, log)
			vReader, vProbeInfo, vErr = aReader, aProbeInfo, aErr
			if aReader != nil {
				closeOnDone = append(closeOnDone, aReader)
			}
		}
	} else if aYDLFormat != nil && vYDLFormat == nil {
		// only audio format
		aReader, aProbeInfo, aErr = downloadAndProbeFormat(ctx, ydl, aYDLFormat.FormatID, log)
		if aReader != nil {
			closeOnDone = append(closeOnDone, aReader)
		}
	} else {
		// don't know, download and probe
		aReader, aProbeInfo, aErr = downloadAndProbeFormat(ctx, ydl, "best", log)
		vReader, vProbeInfo, vErr = aReader, aProbeInfo, aErr
		if aReader != nil {
			closeOnDone = append(closeOnDone, aReader)
		}
	}
	if aErr != nil || vErr != nil {
		return nil, fmt.Errorf("failed to probe")
	}

	log.Printf("Stream mapping:")

	var streamMaps []ffmpeg.StreamMap
	ffmpegFormatFlags := outFormat.FormatFlags

	if len(outFormat.ACodecs) > 0 && aProbeInfo != nil && aProbeInfo.ACodec() != "" {
		codecFormat := chooseFormatCodec(outFormat.ACodecs, aProbeInfo.ACodec())
		streamMaps = append(streamMaps, ffmpeg.StreamMap{
			Reader:     aReader,
			Specifier:  "a:0",
			Codec:      "acodec:" + codecFormat.Codec,
			CodecFlags: codecFormat.CodecFlags,
		})
		ffmpegFormatFlags = append(ffmpegFormatFlags, codecFormat.FormatFlags...)

		log.Printf("  audio %s probed:%s -> %s",
			fancyYDLFormatName(aYDLFormat),
			aProbeInfo,
			codecFormat.Codec,
		)
	}
	if len(outFormat.VCodecs) > 0 && vProbeInfo != nil && vProbeInfo.VCodec() != "" {
		codecFormat := chooseFormatCodec(outFormat.VCodecs, vProbeInfo.VCodec())
		streamMaps = append(streamMaps, ffmpeg.StreamMap{
			Reader:     vReader,
			Specifier:  "v:0",
			Codec:      "vcodec:" + codecFormat.Codec,
			CodecFlags: codecFormat.CodecFlags,
		})
		ffmpegFormatFlags = append(ffmpegFormatFlags, codecFormat.FormatFlags...)

		log.Printf("  video %s probed:%s -> %s",
			fancyYDLFormatName(vYDLFormat),
			vProbeInfo,
			codecFormat.Codec,
		)
	}

	var ffmpegStderr io.Writer
	ffmpegStderr = writelogger.New(log, "ffmpeg stderr> ")
	ffmpegR, ffmpegW := io.Pipe()
	closeOnDone = append(closeOnDone, ffmpegR)

	ffmpegP := &ffmpeg.FFmpeg{
		StreamMaps: streamMaps,
		Format:     ffmpeg.Format{Name: outFormat.Formats.first(), Flags: ffmpegFormatFlags},
		DebugLog:   log,
		Stdout:     ffmpegW,
		Stderr:     ffmpegStderr,
	}

	if err := ffmpegP.Start(ctx); err != nil {
		return nil, err
	}

	// goroutine will take care of closing
	deferCloseFn = nil

	var w io.WriteCloser
	dr.Media, w = io.Pipe()
	closeOnDone = append(closeOnDone, w)

	go func() {
		if outFormat.Prepend == "id3v2" {
			id3v2.Write(w, id3v2FramesFromYoutueDLInfo(ydl))
		}
		io.Copy(w, ffmpegR)
		closeOnDoneFn()
		ffmpegP.Wait()
	}()

	return dr, nil
}
<|endoftext|>


<commit_before>/*
 * ZDNS Copyright 2016 Regents of the University of Michigan
 *
 * 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 main

import (
	"flag"
	"os"
	"runtime"
	"strings"
	"time"

	"github.com/miekg/dns"
	log "github.com/sirupsen/logrus"
	"github.com/zmap/zdns"
	_ "github.com/zmap/zdns/modules/alookup"
	_ "github.com/zmap/zdns/modules/axfr"
	_ "github.com/zmap/zdns/modules/dmarc"
	_ "github.com/zmap/zdns/modules/miekg"
	_ "github.com/zmap/zdns/modules/mxlookup"
	_ "github.com/zmap/zdns/modules/nslookup"
	_ "github.com/zmap/zdns/modules/spf"

	_ "github.com/zmap/zdns/iohandlers/file"
)

func main() {

	var gc zdns.GlobalConf
	// global flags relevant to every lookup module
	flags := flag.NewFlagSet("flags", flag.ExitOnError)
	flags.IntVar(&gc.Threads, "threads", 1000, "number of lightweight go threads")
	flags.IntVar(&gc.GoMaxProcs, "go-processes", 0, "number of OS processes (GOMAXPROCS)")
	flags.StringVar(&gc.NamePrefix, "prefix", "", "name to be prepended to what's passed in (e.g., www.)")
	flags.BoolVar(&gc.AlexaFormat, "alexa", false, "is input file from Alexa Top Million download")
	flags.BoolVar(&gc.IterativeResolution, "iterative", false, "Perform own iteration instead of relying on recursive resolver")
	flags.BoolVar(&gc.Trace, "trace", false, "Output a trace of individual steps for each resolution")
	flags.StringVar(&gc.InputFilePath, "input-file", "-", "names to read")
	flags.StringVar(&gc.OutputFilePath, "output-file", "-", "where should JSON output be saved")
	flags.StringVar(&gc.MetadataFilePath, "metadata-file", "", "where should JSON metadata be saved")
	flags.StringVar(&gc.LogFilePath, "log-file", "", "where should JSON logs be saved")
	flags.IntVar(&gc.Verbosity, "verbosity", 3, "log verbosity: 1 (lowest)--5 (highest)")
	flags.IntVar(&gc.Retries, "retries", 1, "how many times should zdns retry query if timeout or temporary failure")
	flags.IntVar(&gc.MaxDepth, "max-depth", 10, "how deep should we recurse when performing iterative lookups")
	flags.IntVar(&gc.CacheSize, "cache-size", 10000, "how many items can be stored in internal recursive cache")
	flags.StringVar(&gc.InputHandler, "input-handler", "file", "handler to input names")
	flags.StringVar(&gc.OutputHandler, "output-handler", "file", "handler to output names")
	flags.BoolVar(&gc.TCPOnly, "tcp-only", false, "Only perform lookups over TCP")
	flags.BoolVar(&gc.UDPOnly, "udp-only", false, "Only perform lookups over UDP")
	servers_string := flags.String("name-servers", "", "comma-delimited list of DNS servers to use")
	config_file := flags.String("conf-file", "/etc/resolv.conf", "config file for DNS servers")
	timeout := flags.Int("timeout", 15, "timeout for resolving an individual name")
	iterationTimeout := flags.Int("iteration-timeout", 4, "timeout for resolving a single iteration in an iterative query")
	class_string := flags.String("class", "INET", "DNS class to query (INET, CSNET, CHAOS, HESIOD, NONE, ANY (default INET)")
	nanoSeconds := flags.Bool("nanoseconds", false, "Use nanosecond resolution timestamps")
	// allow module to initialize and add its own flags before we parse
	if len(os.Args) < 2 {
		log.Fatal("No lookup module specified. Valid modules: ", zdns.ValidlookupsString())
	}
	gc.Module = strings.ToUpper(os.Args[1])
	factory := zdns.GetLookup(gc.Module)
	if factory == nil {
		flags.Parse(os.Args[1:])
		log.Fatal("Invalid lookup module specified. Valid modules: ", zdns.ValidlookupsString())
	}
	factory.AddFlags(flags)
	flags.Parse(os.Args[2:])
	// Do some basic sanity checking
	// setup global logging
	if gc.LogFilePath != "" {
		f, err := os.OpenFile(gc.LogFilePath, os.O_WRONLY|os.O_CREATE, 0666)
		if err != nil {
			log.Fatalf("Unable to open log file (%s): %s", gc.LogFilePath, err.Error())
		}
		log.SetOutput(f)
	}
	// Translate the assigned verbosity level to a logrus log level.
	switch gc.Verbosity {
	case 1: // Fatal
		log.SetLevel(log.FatalLevel)
	case 2: // Error
		log.SetLevel(log.ErrorLevel)
	case 3: // Warnings  (default)
		log.SetLevel(log.WarnLevel)
	case 4: // Information
		log.SetLevel(log.InfoLevel)
	case 5: // Debugging
		log.SetLevel(log.DebugLevel)
	default:
		log.Fatal("Unknown verbosity level specified. Must be between 1 (lowest)--5 (highest)")
	}
	// complete post facto global initialization based on command line arguments
	gc.Timeout = time.Duration(time.Second * time.Duration(*timeout))
	gc.IterationTimeout = time.Duration(time.Second * time.Duration(*iterationTimeout))
	// class initialization
	switch strings.ToUpper(*class_string) {
	case "INET", "IN":
		gc.Class = dns.ClassINET
	case "CSNET", "CS":
		gc.Class = dns.ClassCSNET
	case "CHAOS", "CH":
		gc.Class = dns.ClassCHAOS
	case "HESIOD", "HS":
		gc.Class = dns.ClassHESIOD
	case "NONE":
		gc.Class = dns.ClassNONE
	case "ANY":
		gc.Class = dns.ClassANY
	default:
		log.Fatal("Unknown record class specified. Valid valued are INET (default), CSNET, CHAOS, HESIOD, NONE, ANY")

	}
	if *servers_string == "" {
		// if we're doing recursive resolution, figure out default OS name servers
		// otherwise, use the set of 13 root name servers
		if gc.IterativeResolution {
			gc.NameServers = zdns.RootServers[:]
		} else {
			ns, err := zdns.GetDNSServers(*config_file)
			if err != nil {
				log.Fatal("Unable to fetch correct name servers:", err.Error())
			}
			gc.NameServers = ns
		}
		gc.NameServersSpecified = false
		log.Info("no name servers specified. will use: ", strings.Join(gc.NameServers, ", "))
	} else {
		gc.NameServers = strings.Split(*servers_string, ",")
		gc.NameServersSpecified = true
	}
	if *nanoSeconds {
		gc.TimeFormat = time.RFC3339Nano
	} else {
		gc.TimeFormat = time.RFC3339
	}
	if gc.GoMaxProcs < 0 {
		log.Fatal("Invalid argument for --go-processes. Must be >1.")
	}
	if gc.GoMaxProcs != 0 {
		runtime.GOMAXPROCS(gc.GoMaxProcs)
	}
	if gc.UDPOnly && gc.TCPOnly {
		log.Fatal("TCP Only and UDP Only are conflicting")
	}
	if len(flags.Args()) > 0 {
		stat, _ := os.Stdin.Stat()
		// If stdin is piped from the terminal, and we havent specified a file, and if we have unparsed args
		// use them for a dig like reslution
		if (stat.Mode()&os.ModeCharDevice) != 0 && gc.InputFilePath == "-" && len(flags.Args()) == 1 {
			gc.PassedName = flags.Args()[0]
		} else {
			log.Fatal("Unused command line flags: ", flags.Args())
		}
	}

	// some modules require multiple passes over a file (this is really just the case for zone files)
	if !factory.AllowStdIn() && gc.InputFilePath == "-" {
		log.Fatal("Specified module does not allow reading from stdin")
	}

	// allow the factory to initialize itself
	if err := factory.Initialize(&gc); err != nil {
		log.Fatal("Factory was unable to initialize:", err.Error())
	}
	// run it.
	if err := zdns.DoLookups(&factory, &gc); err != nil {
		log.Fatal("Unable to run lookups:", err.Error())
	}
	// allow the factory to initialize itself
	if err := factory.Finalize(); err != nil {
		log.Fatal("Factory was unable to finalize:", err.Error())
	}
}
<commit_msg>Default dns server port (#174)<commit_after>/*
 * ZDNS Copyright 2016 Regents of the University of Michigan
 *
 * 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 main

import (
	"flag"
	"os"
	"runtime"
	"strings"
	"time"

	"github.com/miekg/dns"
	log "github.com/sirupsen/logrus"
	"github.com/zmap/zdns"
	_ "github.com/zmap/zdns/modules/alookup"
	_ "github.com/zmap/zdns/modules/axfr"
	_ "github.com/zmap/zdns/modules/dmarc"
	_ "github.com/zmap/zdns/modules/miekg"
	_ "github.com/zmap/zdns/modules/mxlookup"
	_ "github.com/zmap/zdns/modules/nslookup"
	_ "github.com/zmap/zdns/modules/spf"

	_ "github.com/zmap/zdns/iohandlers/file"
)

func main() {

	var gc zdns.GlobalConf
	// global flags relevant to every lookup module
	flags := flag.NewFlagSet("flags", flag.ExitOnError)
	flags.IntVar(&gc.Threads, "threads", 1000, "number of lightweight go threads")
	flags.IntVar(&gc.GoMaxProcs, "go-processes", 0, "number of OS processes (GOMAXPROCS)")
	flags.StringVar(&gc.NamePrefix, "prefix", "", "name to be prepended to what's passed in (e.g., www.)")
	flags.BoolVar(&gc.AlexaFormat, "alexa", false, "is input file from Alexa Top Million download")
	flags.BoolVar(&gc.IterativeResolution, "iterative", false, "Perform own iteration instead of relying on recursive resolver")
	flags.BoolVar(&gc.Trace, "trace", false, "Output a trace of individual steps for each resolution")
	flags.StringVar(&gc.InputFilePath, "input-file", "-", "names to read")
	flags.StringVar(&gc.OutputFilePath, "output-file", "-", "where should JSON output be saved")
	flags.StringVar(&gc.MetadataFilePath, "metadata-file", "", "where should JSON metadata be saved")
	flags.StringVar(&gc.LogFilePath, "log-file", "", "where should JSON logs be saved")
	flags.IntVar(&gc.Verbosity, "verbosity", 3, "log verbosity: 1 (lowest)--5 (highest)")
	flags.IntVar(&gc.Retries, "retries", 1, "how many times should zdns retry query if timeout or temporary failure")
	flags.IntVar(&gc.MaxDepth, "max-depth", 10, "how deep should we recurse when performing iterative lookups")
	flags.IntVar(&gc.CacheSize, "cache-size", 10000, "how many items can be stored in internal recursive cache")
	flags.StringVar(&gc.InputHandler, "input-handler", "file", "handler to input names")
	flags.StringVar(&gc.OutputHandler, "output-handler", "file", "handler to output names")
	flags.BoolVar(&gc.TCPOnly, "tcp-only", false, "Only perform lookups over TCP")
	flags.BoolVar(&gc.UDPOnly, "udp-only", false, "Only perform lookups over UDP")
	servers_string := flags.String("name-servers", "", "comma-delimited list of DNS servers to use. If no port is specified, defaults to :53.")
	config_file := flags.String("conf-file", "/etc/resolv.conf", "config file for DNS servers")
	timeout := flags.Int("timeout", 15, "timeout for resolving an individual name")
	iterationTimeout := flags.Int("iteration-timeout", 4, "timeout for resolving a single iteration in an iterative query")
	class_string := flags.String("class", "INET", "DNS class to query (INET, CSNET, CHAOS, HESIOD, NONE, ANY (default INET)")
	nanoSeconds := flags.Bool("nanoseconds", false, "Use nanosecond resolution timestamps")
	// allow module to initialize and add its own flags before we parse
	if len(os.Args) < 2 {
		log.Fatal("No lookup module specified. Valid modules: ", zdns.ValidlookupsString())
	}
	gc.Module = strings.ToUpper(os.Args[1])
	factory := zdns.GetLookup(gc.Module)
	if factory == nil {
		flags.Parse(os.Args[1:])
		log.Fatal("Invalid lookup module specified. Valid modules: ", zdns.ValidlookupsString())
	}
	factory.AddFlags(flags)
	flags.Parse(os.Args[2:])
	// Do some basic sanity checking
	// setup global logging
	if gc.LogFilePath != "" {
		f, err := os.OpenFile(gc.LogFilePath, os.O_WRONLY|os.O_CREATE, 0666)
		if err != nil {
			log.Fatalf("Unable to open log file (%s): %s", gc.LogFilePath, err.Error())
		}
		log.SetOutput(f)
	}
	// Translate the assigned verbosity level to a logrus log level.
	switch gc.Verbosity {
	case 1: // Fatal
		log.SetLevel(log.FatalLevel)
	case 2: // Error
		log.SetLevel(log.ErrorLevel)
	case 3: // Warnings  (default)
		log.SetLevel(log.WarnLevel)
	case 4: // Information
		log.SetLevel(log.InfoLevel)
	case 5: // Debugging
		log.SetLevel(log.DebugLevel)
	default:
		log.Fatal("Unknown verbosity level specified. Must be between 1 (lowest)--5 (highest)")
	}
	// complete post facto global initialization based on command line arguments
	gc.Timeout = time.Duration(time.Second * time.Duration(*timeout))
	gc.IterationTimeout = time.Duration(time.Second * time.Duration(*iterationTimeout))
	// class initialization
	switch strings.ToUpper(*class_string) {
	case "INET", "IN":
		gc.Class = dns.ClassINET
	case "CSNET", "CS":
		gc.Class = dns.ClassCSNET
	case "CHAOS", "CH":
		gc.Class = dns.ClassCHAOS
	case "HESIOD", "HS":
		gc.Class = dns.ClassHESIOD
	case "NONE":
		gc.Class = dns.ClassNONE
	case "ANY":
		gc.Class = dns.ClassANY
	default:
		log.Fatal("Unknown record class specified. Valid valued are INET (default), CSNET, CHAOS, HESIOD, NONE, ANY")

	}
	if *servers_string == "" {
		// if we're doing recursive resolution, figure out default OS name servers
		// otherwise, use the set of 13 root name servers
		if gc.IterativeResolution {
			gc.NameServers = zdns.RootServers[:]
		} else {
			ns, err := zdns.GetDNSServers(*config_file)
			if err != nil {
				log.Fatal("Unable to fetch correct name servers:", err.Error())
			}
			gc.NameServers = ns
		}
		gc.NameServersSpecified = false
		log.Info("no name servers specified. will use: ", strings.Join(gc.NameServers, ", "))
	} else {
		ns := strings.Split(*servers_string, ",")
		for i, s := range ns {
			if !strings.Contains(s, ":") {
				ns[i] = s + ":53"
			}
		}
		gc.NameServers = ns
		gc.NameServersSpecified = true
	}
	if *nanoSeconds {
		gc.TimeFormat = time.RFC3339Nano
	} else {
		gc.TimeFormat = time.RFC3339
	}
	if gc.GoMaxProcs < 0 {
		log.Fatal("Invalid argument for --go-processes. Must be >1.")
	}
	if gc.GoMaxProcs != 0 {
		runtime.GOMAXPROCS(gc.GoMaxProcs)
	}
	if gc.UDPOnly && gc.TCPOnly {
		log.Fatal("TCP Only and UDP Only are conflicting")
	}
	if len(flags.Args()) > 0 {
		stat, _ := os.Stdin.Stat()
		// If stdin is piped from the terminal, and we havent specified a file, and if we have unparsed args
		// use them for a dig like reslution
		if (stat.Mode()&os.ModeCharDevice) != 0 && gc.InputFilePath == "-" && len(flags.Args()) == 1 {
			gc.PassedName = flags.Args()[0]
		} else {
			log.Fatal("Unused command line flags: ", flags.Args())
		}
	}

	// some modules require multiple passes over a file (this is really just the case for zone files)
	if !factory.AllowStdIn() && gc.InputFilePath == "-" {
		log.Fatal("Specified module does not allow reading from stdin")
	}

	// allow the factory to initialize itself
	if err := factory.Initialize(&gc); err != nil {
		log.Fatal("Factory was unable to initialize:", err.Error())
	}
	// run it.
	if err := zdns.DoLookups(&factory, &gc); err != nil {
		log.Fatal("Unable to run lookups:", err.Error())
	}
	// allow the factory to initialize itself
	if err := factory.Finalize(); err != nil {
		log.Fatal("Factory was unable to finalize:", err.Error())
	}
}
<|endoftext|>


<commit_before>package main

import (
	"sync"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/common/log"
	"gopkg.in/buildkite/go-buildkite.v2/buildkite"
)

type Exporter struct {
	mutex        sync.RWMutex
	client       *buildkite.Client
	duration     prometheus.Gauge
	error        prometheus.Gauge
	totalScrapes prometheus.Counter
	scrapeErrors *prometheus.CounterVec
	builds       *prometheus.GaugeVec
	agents       prometheus.Gauge
	up           prometheus.Gauge
}

func NewExporter(timeout time.Duration) (*Exporter, error) {
	config, err := buildkite.NewTokenConfig(*buildkiteToken, false)

	if err != nil {
		log.Fatalln("buildkite token failed: %s", err)
	}

	httpclient := config.Client()
	httpclient.Timeout = timeout

	client := buildkite.NewClient(httpclient)

	return &Exporter{
		client: client,
		duration: prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: namespace,
			Name:      "last_scrape_duration_seconds",
			Help:      "Duration of the last scrape of metrics from Buildkite.",
		}),
		totalScrapes: prometheus.NewCounter(prometheus.CounterOpts{
			Namespace: namespace,
			Name:      "scrapes_total",
			Help:      "Total number of times Buildkite was scraped for metrics.",
		}),
		scrapeErrors: prometheus.NewCounterVec(prometheus.CounterOpts{
			Namespace: namespace,
			Name:      "scrape_errors_total",
			Help:      "Total number of times an error occurred scraping Buildkite.",
		}, []string{"collector"}),
		error: prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: namespace,
			Name:      "last_scrape_error",
			Help:      "Whether the last scrape of metrics from Buildkite resulted in an error (1 for error, 0 for success).",
		}),
		up: prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: namespace,
			Name:      "up",
			Help:      "Whether the Buildkite is up.",
		}),
		builds: prometheus.NewGaugeVec(prometheus.GaugeOpts{
			Namespace: namespace,
			Name:      "builds",
			Help:      "Number of Buids",
		}, []string{"state"},
		),
		agents: prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: namespace,
			Name:      "agents",
			Help:      "Number of Agents",
		}),
	}, nil
}

// Describe implements prometheus.Collector.
func (e *Exporter) Describe(ch chan<- *prometheus.Desc) {
	ch <- e.duration.Desc()
	ch <- e.error.Desc()
	ch <- e.totalScrapes.Desc()
	ch <- e.up.Desc()
	ch <- e.agents.Desc()

	e.scrapeErrors.Describe(ch)
	e.builds.Describe(ch)
}

func (e *Exporter) Collect(ch chan<- prometheus.Metric) {
	e.mutex.Lock() // To protect metrics from concurrent collects.
	defer e.mutex.Unlock()

	// Reset GaguesVec
	e.builds.Reset()

	e.scrape()

	e.builds.Collect(ch)
	e.scrapeErrors.Collect(ch)

	ch <- e.agents
	ch <- e.duration
	ch <- e.totalScrapes
	ch <- e.error
	ch <- e.up
}

func (e *Exporter) scrape() {
	var err error
	e.totalScrapes.Inc()
	defer func(begun time.Time) {
		e.duration.Set(time.Since(begun).Seconds())
		if err == nil {
			e.error.Set(0)
		} else {
			e.error.Set(1)
		}
	}(time.Now())

	err = e.scrapeBuilds()
	e.scrapeAgents()
	e.up.Set(1)
}

func (e *Exporter) scrapeBuilds() error {
	states := []string{"running", "scheduled", "passed", "failed", "blocked", "canceled",
		"canceling", "skipped", "not_run", "finished"}

	for _, state := range states {
		e.builds.WithLabelValues(state).Set(0)
	}

	options := &buildkite.BuildsListOptions{}

	builds, response, err := e.client.Builds.ListByOrg(*buildkiteOrgName, options)
	if err != nil {
		e.scrapeErrors.WithLabelValues("builds").Inc()
		log.With("collector", "builds").Errorln(err)
		return err
	}
	defer response.Body.Close()

	for _, build := range builds {
		e.builds.With(prometheus.Labels{"state": *build.State}).Inc()
	}

	return nil
}

func (e *Exporter) scrapeAgents() error {
	options := &buildkite.AgentListOptions{}

	agents, response, err := e.client.Agents.List(*buildkiteOrgName, options)
	if err != nil {
		e.scrapeErrors.WithLabelValues("agents").Inc()
		log.With("collector", "agents").Errorln(err)
		return err
	}
	defer response.Body.Close()

	num_agents := len(agents)
	e.agents.Set(float64(num_agents))

	return nil
}
<commit_msg>golint<commit_after>package main

import (
	"sync"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/common/log"
	"gopkg.in/buildkite/go-buildkite.v2/buildkite"
)

// Exporter collects Buildkite metrics. It implements prometheus.Collector.
type Exporter struct {
	mutex        sync.RWMutex
	client       *buildkite.Client
	duration     prometheus.Gauge
	error        prometheus.Gauge
	totalScrapes prometheus.Counter
	scrapeErrors *prometheus.CounterVec
	builds       *prometheus.GaugeVec
	agents       prometheus.Gauge
	up           prometheus.Gauge
}

// NewExporter returns a new BuildKite exporter
func NewExporter(timeout time.Duration) (*Exporter, error) {
	config, err := buildkite.NewTokenConfig(*buildkiteToken, false)

	if err != nil {
		log.Fatalln("buildkite token failed: %s", err)
	}

	httpclient := config.Client()
	httpclient.Timeout = timeout

	client := buildkite.NewClient(httpclient)

	return &Exporter{
		client: client,
		duration: prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: namespace,
			Name:      "last_scrape_duration_seconds",
			Help:      "Duration of the last scrape of metrics from Buildkite.",
		}),
		totalScrapes: prometheus.NewCounter(prometheus.CounterOpts{
			Namespace: namespace,
			Name:      "scrapes_total",
			Help:      "Total number of times Buildkite was scraped for metrics.",
		}),
		scrapeErrors: prometheus.NewCounterVec(prometheus.CounterOpts{
			Namespace: namespace,
			Name:      "scrape_errors_total",
			Help:      "Total number of times an error occurred scraping Buildkite.",
		}, []string{"collector"}),
		error: prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: namespace,
			Name:      "last_scrape_error",
			Help:      "Whether the last scrape of metrics from Buildkite resulted in an error (1 for error, 0 for success).",
		}),
		up: prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: namespace,
			Name:      "up",
			Help:      "Whether the Buildkite is up.",
		}),
		builds: prometheus.NewGaugeVec(prometheus.GaugeOpts{
			Namespace: namespace,
			Name:      "builds",
			Help:      "Number of Buids",
		}, []string{"state"},
		),
		agents: prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: namespace,
			Name:      "agents",
			Help:      "Number of Agents",
		}),
	}, nil
}

// Describe implements prometheus.Collector.
func (e *Exporter) Describe(ch chan<- *prometheus.Desc) {
	ch <- e.duration.Desc()
	ch <- e.error.Desc()
	ch <- e.totalScrapes.Desc()
	ch <- e.up.Desc()
	ch <- e.agents.Desc()

	e.scrapeErrors.Describe(ch)
	e.builds.Describe(ch)
}

// Collect implements prometheus.Collector.
func (e *Exporter) Collect(ch chan<- prometheus.Metric) {
	e.mutex.Lock() // To protect metrics from concurrent collects.
	defer e.mutex.Unlock()

	// Reset GaguesVec
	e.builds.Reset()

	e.scrape()

	e.builds.Collect(ch)
	e.scrapeErrors.Collect(ch)

	ch <- e.agents
	ch <- e.duration
	ch <- e.totalScrapes
	ch <- e.error
	ch <- e.up
}

func (e *Exporter) scrape() {
	var err error
	e.totalScrapes.Inc()
	defer func(begun time.Time) {
		e.duration.Set(time.Since(begun).Seconds())
		if err == nil {
			e.error.Set(0)
		} else {
			e.error.Set(1)
		}
	}(time.Now())

	err = e.scrapeBuilds()
	e.scrapeAgents()
	e.up.Set(1)
}

func (e *Exporter) scrapeBuilds() error {
	states := []string{"running", "scheduled", "passed", "failed", "blocked", "canceled",
		"canceling", "skipped", "not_run", "finished"}

	for _, state := range states {
		e.builds.WithLabelValues(state).Set(0)
	}

	options := &buildkite.BuildsListOptions{}

	builds, response, err := e.client.Builds.ListByOrg(*buildkiteOrgName, options)
	if err != nil {
		e.scrapeErrors.WithLabelValues("builds").Inc()
		log.With("collector", "builds").Errorln(err)
		return err
	}
	defer response.Body.Close()

	for _, build := range builds {
		e.builds.With(prometheus.Labels{"state": *build.State}).Inc()
	}

	return nil
}

func (e *Exporter) scrapeAgents() error {
	options := &buildkite.AgentListOptions{}

	agents, response, err := e.client.Agents.List(*buildkiteOrgName, options)
	if err != nil {
		e.scrapeErrors.WithLabelValues("agents").Inc()
		log.With("collector", "agents").Errorln(err)
		return err
	}
	defer response.Body.Close()

	numAgents := len(agents)
	e.agents.Set(float64(numAgents))

	return nil
}
<|endoftext|>


<commit_before>package Common

// 每日滚动的LOG实现
import (
	"fmt"
	"os"
	"path/filepath"
	"runtime"
	"time"
)

const (
	LogLevelDrop = iota
	LogLevelDebug
	LogLevelInfo
	LogLevelWarn
	LogLevelError
	LogLevelNone
)

var (
	logDir    = ""
	logDay    = 0
	logLevel  = LogLevelInfo
	logFile   = os.Stderr
	logLock   = NewLock()
	logTicker = time.NewTicker(time.Second)
	logSlice  = make([]interface{}, 1, 1024)
)

func SetLogDir(dir string) {
	logDir = dir
}

func SetLogLevel(level int) {
	logLevel = level
}

func check() {
	select {
	case <-logTicker.C:
	default:
		return
	}

	if len(logDir) == 0 {
		return
	}

	if logLock.TryLock() {
		defer logLock.Unlock()
	} else {
		return
	}

	now := time.Now().UTC()
	if logDay == now.Day() {
		return
	}

	if logFile != os.Stderr {
		logFile.Close()
	}

	logDay = now.Day()
	logProc := filepath.Base(os.Args[0])
	filename := filepath.Join(logDir, fmt.Sprintf("%s.%s.log", logProc, now.Format("2006-01-02")))

	newlog, err := os.OpenFile(filename, os.O_WRONLY|os.O_APPEND|os.O_CREATE, 0666)
	if err != nil {
		fmt.Fprintln(os.Stderr, NumberUTC(), "open log file", filename, err, "use STDOUT")
		logFile = os.Stderr
	} else {
		logFile = newlog
	}
}

func filebase(file string) string {
	beg, end := len(file)-1, len(file)
	for ; beg >= 0; beg-- {
		if os.IsPathSeparator(file[beg]) {
			beg++
			break
		} else if file[beg] == '.' {
			end = beg
		}
	}
	return file[beg:end]
}

func prefix(level string) {
	_, file, line, _ := runtime.Caller(2)
	logSlice[0] = fmt.Sprintf("%d %s %s[%d]:", NumberUTC(), level, filebase(file), line)
}

func DropLog(v ...interface{}) {}

func DebugLog(v ...interface{}) {
	if logLevel > LogLevelDebug {
		return
	}

	check()
	prefix("debug")

	out := append(logSlice, v...)
	logLock.Lock()
	defer logLock.Unlock()
	fmt.Fprintln(logFile, out...)
}

func InfoLog(v ...interface{}) {
	if logLevel > LogLevelInfo {
		return
	}

	check()
	prefix("info")

	out := append(logSlice, v...)
	logLock.Lock()
	defer logLock.Unlock()
	fmt.Fprintln(logFile, out...)
}

func WarningLog(v ...interface{}) {
	if logLevel > LogLevelWarn {
		return
	}

	check()
	prefix("warn")

	out := append(logSlice, v...)
	logLock.Lock()
	defer logLock.Unlock()
	fmt.Fprintln(logFile, out...)
}

func ErrorLog(v ...interface{}) {
	if logLevel > LogLevelError {
		return
	}

	check()
	prefix("error")

	out := append(logSlice, v...)
	logLock.Lock()
	defer logLock.Unlock()
	fmt.Fprintln(logFile, out...)
}
<commit_msg>fix init log bug<commit_after>package Common

// 每日滚动的LOG实现
import (
	"fmt"
	"os"
	"path/filepath"
	"runtime"
	"time"
)

const (
	LogLevelDrop = iota
	LogLevelDebug
	LogLevelInfo
	LogLevelWarn
	LogLevelError
	LogLevelNone
)

var (
	logDir    = ""
	logDay    = 0
	logLevel  = LogLevelInfo
	logFile   = os.Stderr
	logLock   = NewLock()
	logTicker = time.NewTicker(time.Second)
	logSlice  = make([]interface{}, 1, 1024)
)

func init() {
	select {
	case logTicker <- time.Now():
	default:
	}
}

func SetLogDir(dir string) {
	logDir = dir
}

func SetLogLevel(level int) {
	logLevel = level
}

func check() {
	select {
	case <-logTicker.C:
	default:
		return
	}

	if len(logDir) == 0 {
		return
	}

	if logLock.TryLock() {
		defer logLock.Unlock()
	} else {
		return
	}

	now := time.Now().UTC()
	if logDay == now.Day() {
		return
	}

	if logFile != os.Stderr {
		logFile.Close()
	}

	logDay = now.Day()
	logProc := filepath.Base(os.Args[0])
	filename := filepath.Join(logDir, fmt.Sprintf("%s.%s.log", logProc, now.Format("2006-01-02")))

	newlog, err := os.OpenFile(filename, os.O_WRONLY|os.O_APPEND|os.O_CREATE, 0666)
	if err != nil {
		fmt.Fprintln(os.Stderr, NumberUTC(), "open log file", filename, err, "use STDOUT")
		logFile = os.Stderr
	} else {
		logFile = newlog
	}
}

func filebase(file string) string {
	beg, end := len(file)-1, len(file)
	for ; beg >= 0; beg-- {
		if os.IsPathSeparator(file[beg]) {
			beg++
			break
		} else if file[beg] == '.' {
			end = beg
		}
	}
	return file[beg:end]
}

func prefix(level string) {
	_, file, line, _ := runtime.Caller(2)
	logSlice[0] = fmt.Sprintf("%d %s %s[%d]:", NumberUTC(), level, filebase(file), line)
}

func DropLog(v ...interface{}) {}

func DebugLog(v ...interface{}) {
	if logLevel > LogLevelDebug {
		return
	}

	check()
	prefix("debug")

	out := append(logSlice, v...)
	logLock.Lock()
	defer logLock.Unlock()
	fmt.Fprintln(logFile, out...)
}

func InfoLog(v ...interface{}) {
	if logLevel > LogLevelInfo {
		return
	}

	check()
	prefix("info")

	out := append(logSlice, v...)
	logLock.Lock()
	defer logLock.Unlock()
	fmt.Fprintln(logFile, out...)
}

func WarningLog(v ...interface{}) {
	if logLevel > LogLevelWarn {
		return
	}

	check()
	prefix("warn")

	out := append(logSlice, v...)
	logLock.Lock()
	defer logLock.Unlock()
	fmt.Fprintln(logFile, out...)
}

func ErrorLog(v ...interface{}) {
	if logLevel > LogLevelError {
		return
	}

	check()
	prefix("error")

	out := append(logSlice, v...)
	logLock.Lock()
	defer logLock.Unlock()
	fmt.Fprintln(logFile, out...)
}
<|endoftext|>


<commit_before>package arn

import (
	"errors"
	"fmt"
	"reflect"

	"github.com/aerogo/aero"
	"github.com/aerogo/api"
)

// Force interface implementations
var (
	_ Publishable   = (*Company)(nil)
	_ fmt.Stringer  = (*Company)(nil)
	_ api.Newable   = (*Company)(nil)
	_ api.Editable  = (*Company)(nil)
	_ api.Deletable = (*Company)(nil)
)

// Actions
func init() {
	API.RegisterActions("Company", []*api.Action{
		// Publish
		PublishAction(),

		// Unpublish
		UnpublishAction(),

		// Like
		LikeAction(),

		// Unlike
		UnlikeAction(),
	})
}

// Create sets the data for a new company with data we received from the API request.
func (company *Company) Create(ctx *aero.Context) error {
	user := GetUserFromContext(ctx)

	if user == nil {
		return errors.New("Not logged in")
	}

	company.ID = GenerateID("Company")
	company.Created = DateTimeUTC()
	company.CreatedBy = user.ID
	company.Mappings = []*Mapping{}
	company.Links = []*Link{}
	company.Tags = []string{}
	company.Likes = []string{}

	// Write log entry
	logEntry := NewEditLogEntry(user.ID, "create", "Company", company.ID, "", "", "")
	logEntry.Save()

	return company.Unpublish()
}

// Edit saves a log entry for the edit.
func (company *Company) Edit(ctx *aero.Context, key string, value reflect.Value, newValue reflect.Value) (bool, error) {
	user := GetUserFromContext(ctx)

	// Write log entry
	logEntry := NewEditLogEntry(user.ID, "edit", "Company", company.ID, key, fmt.Sprint(value.Interface()), fmt.Sprint(newValue.Interface()))
	logEntry.Save()

	return false, nil
}

// AfterEdit updates the metadata.
func (company *Company) AfterEdit(ctx *aero.Context) error {
	company.Edited = DateTimeUTC()
	company.EditedBy = GetUserFromContext(ctx).ID
	return nil
}

// Save saves the company in the database.
func (company *Company) Save() {
	DB.Set("Company", company.ID, company)
}

// DeleteInContext deletes the company in the given context.
func (company *Company) DeleteInContext(ctx *aero.Context) error {
	user := GetUserFromContext(ctx)

	// Write log entry
	logEntry := NewEditLogEntry(user.ID, "delete", "Company", company.ID, "", fmt.Sprint(company), "")
	logEntry.Save()

	return company.Delete()
}

// Delete deletes the object from the database.
func (company *Company) Delete() error {
	if company.IsDraft {
		draftIndex := company.Creator().DraftIndex()
		draftIndex.CompanyID = ""
		draftIndex.Save()
	}

	// Remove company ID from all anime
	for anime := range StreamAnime() {
		for index, id := range anime.StudioIDs {
			if id == company.ID {
				anime.StudioIDs = append(anime.StudioIDs[:index], anime.StudioIDs[index+1:]...)
				break
			}
		}

		for index, id := range anime.ProducerIDs {
			if id == company.ID {
				anime.ProducerIDs = append(anime.ProducerIDs[:index], anime.ProducerIDs[index+1:]...)
				break
			}
		}

		for index, id := range anime.LicensorIDs {
			if id == company.ID {
				anime.LicensorIDs = append(anime.LicensorIDs[:index], anime.LicensorIDs[index+1:]...)
				break
			}
		}
	}

	DB.Delete("Company", company.ID)
	return nil
}

// Authorize returns an error if the given API request is not authorized.
func (company *Company) Authorize(ctx *aero.Context, action string) error {
	user := GetUserFromContext(ctx)

	if user == nil {
		return errors.New("Not logged in")
	}

	if user.Role != "editor" && user.Role != "admin" {
		return errors.New("Insufficient permissions")
	}

	return nil
}
<commit_msg>Fixed company event listeners<commit_after>package arn

import (
	"errors"
	"fmt"
	"reflect"

	"github.com/aerogo/aero"
	"github.com/aerogo/api"
)

// Force interface implementations
var (
	_ Publishable            = (*Company)(nil)
	_ fmt.Stringer           = (*Company)(nil)
	_ api.Newable            = (*Company)(nil)
	_ api.Editable           = (*Company)(nil)
	_ api.Deletable          = (*Company)(nil)
	_ api.ArrayEventListener = (*Company)(nil)
)

// Actions
func init() {
	API.RegisterActions("Company", []*api.Action{
		// Publish
		PublishAction(),

		// Unpublish
		UnpublishAction(),

		// Like
		LikeAction(),

		// Unlike
		UnlikeAction(),
	})
}

// Create sets the data for a new company with data we received from the API request.
func (company *Company) Create(ctx *aero.Context) error {
	user := GetUserFromContext(ctx)

	if user == nil {
		return errors.New("Not logged in")
	}

	company.ID = GenerateID("Company")
	company.Created = DateTimeUTC()
	company.CreatedBy = user.ID
	company.Mappings = []*Mapping{}
	company.Links = []*Link{}
	company.Tags = []string{}
	company.Likes = []string{}

	// Write log entry
	logEntry := NewEditLogEntry(user.ID, "create", "Company", company.ID, "", "", "")
	logEntry.Save()

	return company.Unpublish()
}

// Edit saves a log entry for the edit.
func (company *Company) Edit(ctx *aero.Context, key string, value reflect.Value, newValue reflect.Value) (bool, error) {
	user := GetUserFromContext(ctx)

	// Write log entry
	logEntry := NewEditLogEntry(user.ID, "edit", "Company", company.ID, key, fmt.Sprint(value.Interface()), fmt.Sprint(newValue.Interface()))
	logEntry.Save()

	return false, nil
}

// OnAppend saves a log entry.
func (company *Company) OnAppend(ctx *aero.Context, key string, index int, obj interface{}) {
	user := GetUserFromContext(ctx)
	logEntry := NewEditLogEntry(user.ID, "arrayAppend", "Company", company.ID, fmt.Sprintf("%s[%d]", key, index), "", fmt.Sprint(obj))
	logEntry.Save()
}

// OnRemove saves a log entry.
func (company *Company) OnRemove(ctx *aero.Context, key string, index int, obj interface{}) {
	user := GetUserFromContext(ctx)
	logEntry := NewEditLogEntry(user.ID, "arrayRemove", "Company", company.ID, fmt.Sprintf("%s[%d]", key, index), fmt.Sprint(obj), "")
	logEntry.Save()
}

// AfterEdit updates the metadata.
func (company *Company) AfterEdit(ctx *aero.Context) error {
	company.Edited = DateTimeUTC()
	company.EditedBy = GetUserFromContext(ctx).ID
	return nil
}

// Save saves the company in the database.
func (company *Company) Save() {
	DB.Set("Company", company.ID, company)
}

// DeleteInContext deletes the company in the given context.
func (company *Company) DeleteInContext(ctx *aero.Context) error {
	user := GetUserFromContext(ctx)

	// Write log entry
	logEntry := NewEditLogEntry(user.ID, "delete", "Company", company.ID, "", fmt.Sprint(company), "")
	logEntry.Save()

	return company.Delete()
}

// Delete deletes the object from the database.
func (company *Company) Delete() error {
	if company.IsDraft {
		draftIndex := company.Creator().DraftIndex()
		draftIndex.CompanyID = ""
		draftIndex.Save()
	}

	// Remove company ID from all anime
	for anime := range StreamAnime() {
		for index, id := range anime.StudioIDs {
			if id == company.ID {
				anime.StudioIDs = append(anime.StudioIDs[:index], anime.StudioIDs[index+1:]...)
				break
			}
		}

		for index, id := range anime.ProducerIDs {
			if id == company.ID {
				anime.ProducerIDs = append(anime.ProducerIDs[:index], anime.ProducerIDs[index+1:]...)
				break
			}
		}

		for index, id := range anime.LicensorIDs {
			if id == company.ID {
				anime.LicensorIDs = append(anime.LicensorIDs[:index], anime.LicensorIDs[index+1:]...)
				break
			}
		}
	}

	DB.Delete("Company", company.ID)
	return nil
}

// Authorize returns an error if the given API request is not authorized.
func (company *Company) Authorize(ctx *aero.Context, action string) error {
	user := GetUserFromContext(ctx)

	if user == nil {
		return errors.New("Not logged in")
	}

	if user.Role != "editor" && user.Role != "admin" {
		return errors.New("Insufficient permissions")
	}

	return nil
}
<|endoftext|>


<commit_before>package kerberos

// #cgo LDFLAGS: -lkrb5 -lk5crypto -lcom_err
// #include <krb5.h>
import "C"

type Context struct {
	ctx C.krb5_context
}

func NewContext() (*Context, error) {
	var ctx C.krb5_context
	code := C.krb5_init_context(&ctx)
	if code != 0 {
		return nil, Error{
			Code: int(code),
			Msg:  "Cannot initialize Kerberos Context"}
	}
	return &Context{ctx: ctx}, nil
}

func (c *Context) Free() {
	if c.ctx != nil {
		C.krb5_free_context(c.ctx)
		c.ctx = nil
	}
}

func (c *Context) newError(code C.krb5_error_code) error {
	msg := C.krb5_get_error_message(c.ctx, code)
	C.krb5_clear_error_message(c.ctx)
	defer C.krb5_free_error_message(c.ctx, msg)
	return Error{Code: int(code), Msg: C.GoString(msg)}
}

func (c *Context) NewPrincipal(pname string) (*Principal, error) {
	var princ C.krb5_principal
	pnameC := C.CString(pname)
	defer C.krb5_free_string(c.ctx, pnameC)
	code := C.krb5_parse_name(c.ctx, pnameC, &princ)
	if code != 0 {
		return nil, c.newError(code)
	}
	return &Principal{princ: princ}, nil
}

func (c *Context) FreePrincipal(princ *Principal) {
	if princ != nil {
		C.krb5_free_principal(c.ctx, princ.princ)
		princ.princ = nil
	}
}

func (c *Context) Localname(p *Principal) (string, error) {
	var buf *C.char
	defer C.krb5_free_string(c.ctx, buf)
	code := C.krb5_aname_to_localname(c.ctx, p.princ, 255, buf)
	if code != 0 {
		return "", c.newError(code)
	}
	return C.GoString(buf), nil
}

func (c *Context) Unparse(p *Principal) (string, error) {
	var buf *C.char
	defer C.krb5_free_unparsed_name(c.ctx, buf)
	code := C.krb5_unparse_name(c.ctx, p.princ, &buf)
	if code != 0 {
		return "", c.newError(code)
	}
	return C.GoString(buf), nil
}

type Error struct {
	Code int
	Msg  string
}

func (e Error) Error() string {
	return e.Msg
}

type Principal struct {
	princ C.krb5_principal
}
<commit_msg>Well, at least it works now<commit_after>package kerberos

// #cgo LDFLAGS: -lkrb5 -lk5crypto -lcom_err
// #include <krb5.h>
import "C"

type Context struct {
	ctx C.krb5_context
}

func NewContext() (*Context, error) {
	var ctx C.krb5_context
	code := C.krb5_init_context(&ctx)
	if code != 0 {
		return nil, Error{
			Code: int(code),
			Msg:  "Cannot initialize Kerberos Context"}
	}
	return &Context{ctx: ctx}, nil
}

func (c *Context) Free() {
	if c.ctx != nil {
		C.krb5_free_context(c.ctx)
		c.ctx = nil
	}
}

func (c *Context) newError(code C.krb5_error_code) error {
	msg := C.krb5_get_error_message(c.ctx, code)
	C.krb5_clear_error_message(c.ctx)
	defer C.krb5_free_error_message(c.ctx, msg)
	return Error{Code: int(code), Msg: C.GoString(msg)}
}

func (c *Context) NewPrincipal(pname string) (*Principal, error) {
	var princ C.krb5_principal
	pnameC := C.CString(pname)
	defer C.krb5_free_string(c.ctx, pnameC)
	code := C.krb5_parse_name(c.ctx, pnameC, &princ)
	if code != 0 {
		return nil, c.newError(code)
	}
	return &Principal{princ: princ}, nil
}

func (c *Context) FreePrincipal(princ *Principal) {
	if princ != nil {
		C.krb5_free_principal(c.ctx, princ.princ)
		princ.princ = nil
	}
}

func (c *Context) Localname(p *Principal) (string, error) {
	var buf C.char
	// defer C.krb5_free_string(c.ctx, &buf)
	code := C.krb5_aname_to_localname(c.ctx, p.princ, 255, &buf)
	if code != 0 {
		return "", c.newError(code)
	}
	return C.GoString(&buf), nil
}

func (c *Context) Unparse(p *Principal) (string, error) {
	var buf *C.char
	defer C.krb5_free_unparsed_name(c.ctx, buf)
	code := C.krb5_unparse_name(c.ctx, p.princ, &buf)
	if code != 0 {
		return "", c.newError(code)
	}
	return C.GoString(buf), nil
}

type Error struct {
	Code int
	Msg  string
}

func (e Error) Error() string {
	return e.Msg
}

type Principal struct {
	princ C.krb5_principal
}
<|endoftext|>


<commit_before>// Copyright 2017 Google Inc. All Rights Reserved.
//
// 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 firebase is the entry point to the Firebase Admin SDK. It provides functionality for initializing App
// instances, which serve as the central entities that provide access to various other Firebase services exposed
// from the SDK.
package firebase

import (
	"context"
	"encoding/json"
	"errors"
	"io/ioutil"
	"os"

	"cloud.google.com/go/firestore"
	"firebase.google.com/go/v4/auth"
	"firebase.google.com/go/v4/db"
	"firebase.google.com/go/v4/iid"
	"firebase.google.com/go/v4/internal"
	"firebase.google.com/go/v4/messaging"
	"firebase.google.com/go/v4/storage"
	"google.golang.org/api/option"
	"google.golang.org/api/transport"
)

var defaultAuthOverrides = make(map[string]interface{})

// Version of the Firebase Go Admin SDK.
const Version = "4.6.1"

// firebaseEnvName is the name of the environment variable with the Config.
const firebaseEnvName = "FIREBASE_CONFIG"

// An App holds configuration and state common to all Firebase services that are exposed from the SDK.
type App struct {
	authOverride     map[string]interface{}
	dbURL            string
	projectID        string
	serviceAccountID string
	storageBucket    string
	opts             []option.ClientOption
}

// Config represents the configuration used to initialize an App.
type Config struct {
	AuthOverride     *map[string]interface{} `json:"databaseAuthVariableOverride"`
	DatabaseURL      string                  `json:"databaseURL"`
	ProjectID        string                  `json:"projectId"`
	ServiceAccountID string                  `json:"serviceAccountId"`
	StorageBucket    string                  `json:"storageBucket"`
}

// Auth returns an instance of auth.Client.
func (a *App) Auth(ctx context.Context) (*auth.Client, error) {
	conf := &internal.AuthConfig{
		ProjectID:        a.projectID,
		Opts:             a.opts,
		ServiceAccountID: a.serviceAccountID,
		Version:          Version,
	}
	return auth.NewClient(ctx, conf)
}

// Database returns an instance of db.Client to interact with the default Firebase Database
// configured via Config.DatabaseURL.
func (a *App) Database(ctx context.Context) (*db.Client, error) {
	return a.DatabaseWithURL(ctx, a.dbURL)
}

// DatabaseWithURL returns an instance of db.Client to interact with the Firebase Database
// identified by the given URL.
func (a *App) DatabaseWithURL(ctx context.Context, url string) (*db.Client, error) {
	conf := &internal.DatabaseConfig{
		AuthOverride: a.authOverride,
		URL:          url,
		Opts:         a.opts,
		Version:      Version,
	}
	return db.NewClient(ctx, conf)
}

// Storage returns a new instance of storage.Client.
func (a *App) Storage(ctx context.Context) (*storage.Client, error) {
	conf := &internal.StorageConfig{
		Opts:   a.opts,
		Bucket: a.storageBucket,
	}
	return storage.NewClient(ctx, conf)
}

// Firestore returns a new firestore.Client instance from the https://godoc.org/cloud.google.com/go/firestore
// package.
func (a *App) Firestore(ctx context.Context) (*firestore.Client, error) {
	if a.projectID == "" {
		return nil, errors.New("project id is required to access Firestore")
	}
	return firestore.NewClient(ctx, a.projectID, a.opts...)
}

// InstanceID returns an instance of iid.Client.
func (a *App) InstanceID(ctx context.Context) (*iid.Client, error) {
	conf := &internal.InstanceIDConfig{
		ProjectID: a.projectID,
		Opts:      a.opts,
	}
	return iid.NewClient(ctx, conf)
}

// Messaging returns an instance of messaging.Client.
func (a *App) Messaging(ctx context.Context) (*messaging.Client, error) {
	conf := &internal.MessagingConfig{
		ProjectID: a.projectID,
		Opts:      a.opts,
		Version:   Version,
	}
	return messaging.NewClient(ctx, conf)
}

// NewApp creates a new App from the provided config and client options.
//
// If the client options contain a valid credential (a service account file, a refresh token
// file or an oauth2.TokenSource) the App will be authenticated using that credential. Otherwise,
// NewApp attempts to authenticate the App with Google application default credentials.
// If `config` is nil, the SDK will attempt to load the config options from the
// `FIREBASE_CONFIG` environment variable. If the value in it starts with a `{` it is parsed as a
// JSON object, otherwise it is assumed to be the name of the JSON file containing the options.
func NewApp(ctx context.Context, config *Config, opts ...option.ClientOption) (*App, error) {
	o := []option.ClientOption{option.WithScopes(internal.FirebaseScopes...)}
	o = append(o, opts...)
	if config == nil {
		var err error
		if config, err = getConfigDefaults(); err != nil {
			return nil, err
		}
	}

	pid := getProjectID(ctx, config, o...)
	ao := defaultAuthOverrides
	if config.AuthOverride != nil {
		ao = *config.AuthOverride
	}

	return &App{
		authOverride:     ao,
		dbURL:            config.DatabaseURL,
		projectID:        pid,
		serviceAccountID: config.ServiceAccountID,
		storageBucket:    config.StorageBucket,
		opts:             o,
	}, nil
}

// getConfigDefaults reads the default config file, defined by the FIREBASE_CONFIG
// env variable, used only when options are nil.
func getConfigDefaults() (*Config, error) {
	fbc := &Config{}
	confFileName := os.Getenv(firebaseEnvName)
	if confFileName == "" {
		return fbc, nil
	}
	var dat []byte
	if confFileName[0] == byte('{') {
		dat = []byte(confFileName)
	} else {
		var err error
		if dat, err = ioutil.ReadFile(confFileName); err != nil {
			return nil, err
		}
	}
	if err := json.Unmarshal(dat, fbc); err != nil {
		return nil, err
	}

	// Some special handling necessary for db auth overrides
	var m map[string]interface{}
	if err := json.Unmarshal(dat, &m); err != nil {
		return nil, err
	}
	if ao, ok := m["databaseAuthVariableOverride"]; ok && ao == nil {
		// Auth overrides are explicitly set to null
		var nullMap map[string]interface{}
		fbc.AuthOverride = &nullMap
	}
	return fbc, nil
}

func getProjectID(ctx context.Context, config *Config, opts ...option.ClientOption) string {
	if config.ProjectID != "" {
		return config.ProjectID
	}

	creds, _ := transport.Creds(ctx, opts...)
	if creds != nil && creds.ProjectID != "" {
		return creds.ProjectID
	}

	if pid := os.Getenv("GOOGLE_CLOUD_PROJECT"); pid != "" {
		return pid
	}

	return os.Getenv("GCLOUD_PROJECT")
}
<commit_msg>[chore] Release 4.7.0 (#473)<commit_after>// Copyright 2017 Google Inc. All Rights Reserved.
//
// 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 firebase is the entry point to the Firebase Admin SDK. It provides functionality for initializing App
// instances, which serve as the central entities that provide access to various other Firebase services exposed
// from the SDK.
package firebase

import (
	"context"
	"encoding/json"
	"errors"
	"io/ioutil"
	"os"

	"cloud.google.com/go/firestore"
	"firebase.google.com/go/v4/auth"
	"firebase.google.com/go/v4/db"
	"firebase.google.com/go/v4/iid"
	"firebase.google.com/go/v4/internal"
	"firebase.google.com/go/v4/messaging"
	"firebase.google.com/go/v4/storage"
	"google.golang.org/api/option"
	"google.golang.org/api/transport"
)

var defaultAuthOverrides = make(map[string]interface{})

// Version of the Firebase Go Admin SDK.
const Version = "4.7.0"

// firebaseEnvName is the name of the environment variable with the Config.
const firebaseEnvName = "FIREBASE_CONFIG"

// An App holds configuration and state common to all Firebase services that are exposed from the SDK.
type App struct {
	authOverride     map[string]interface{}
	dbURL            string
	projectID        string
	serviceAccountID string
	storageBucket    string
	opts             []option.ClientOption
}

// Config represents the configuration used to initialize an App.
type Config struct {
	AuthOverride     *map[string]interface{} `json:"databaseAuthVariableOverride"`
	DatabaseURL      string                  `json:"databaseURL"`
	ProjectID        string                  `json:"projectId"`
	ServiceAccountID string                  `json:"serviceAccountId"`
	StorageBucket    string                  `json:"storageBucket"`
}

// Auth returns an instance of auth.Client.
func (a *App) Auth(ctx context.Context) (*auth.Client, error) {
	conf := &internal.AuthConfig{
		ProjectID:        a.projectID,
		Opts:             a.opts,
		ServiceAccountID: a.serviceAccountID,
		Version:          Version,
	}
	return auth.NewClient(ctx, conf)
}

// Database returns an instance of db.Client to interact with the default Firebase Database
// configured via Config.DatabaseURL.
func (a *App) Database(ctx context.Context) (*db.Client, error) {
	return a.DatabaseWithURL(ctx, a.dbURL)
}

// DatabaseWithURL returns an instance of db.Client to interact with the Firebase Database
// identified by the given URL.
func (a *App) DatabaseWithURL(ctx context.Context, url string) (*db.Client, error) {
	conf := &internal.DatabaseConfig{
		AuthOverride: a.authOverride,
		URL:          url,
		Opts:         a.opts,
		Version:      Version,
	}
	return db.NewClient(ctx, conf)
}

// Storage returns a new instance of storage.Client.
func (a *App) Storage(ctx context.Context) (*storage.Client, error) {
	conf := &internal.StorageConfig{
		Opts:   a.opts,
		Bucket: a.storageBucket,
	}
	return storage.NewClient(ctx, conf)
}

// Firestore returns a new firestore.Client instance from the https://godoc.org/cloud.google.com/go/firestore
// package.
func (a *App) Firestore(ctx context.Context) (*firestore.Client, error) {
	if a.projectID == "" {
		return nil, errors.New("project id is required to access Firestore")
	}
	return firestore.NewClient(ctx, a.projectID, a.opts...)
}

// InstanceID returns an instance of iid.Client.
func (a *App) InstanceID(ctx context.Context) (*iid.Client, error) {
	conf := &internal.InstanceIDConfig{
		ProjectID: a.projectID,
		Opts:      a.opts,
	}
	return iid.NewClient(ctx, conf)
}

// Messaging returns an instance of messaging.Client.
func (a *App) Messaging(ctx context.Context) (*messaging.Client, error) {
	conf := &internal.MessagingConfig{
		ProjectID: a.projectID,
		Opts:      a.opts,
		Version:   Version,
	}
	return messaging.NewClient(ctx, conf)
}

// NewApp creates a new App from the provided config and client options.
//
// If the client options contain a valid credential (a service account file, a refresh token
// file or an oauth2.TokenSource) the App will be authenticated using that credential. Otherwise,
// NewApp attempts to authenticate the App with Google application default credentials.
// If `config` is nil, the SDK will attempt to load the config options from the
// `FIREBASE_CONFIG` environment variable. If the value in it starts with a `{` it is parsed as a
// JSON object, otherwise it is assumed to be the name of the JSON file containing the options.
func NewApp(ctx context.Context, config *Config, opts ...option.ClientOption) (*App, error) {
	o := []option.ClientOption{option.WithScopes(internal.FirebaseScopes...)}
	o = append(o, opts...)
	if config == nil {
		var err error
		if config, err = getConfigDefaults(); err != nil {
			return nil, err
		}
	}

	pid := getProjectID(ctx, config, o...)
	ao := defaultAuthOverrides
	if config.AuthOverride != nil {
		ao = *config.AuthOverride
	}

	return &App{
		authOverride:     ao,
		dbURL:            config.DatabaseURL,
		projectID:        pid,
		serviceAccountID: config.ServiceAccountID,
		storageBucket:    config.StorageBucket,
		opts:             o,
	}, nil
}

// getConfigDefaults reads the default config file, defined by the FIREBASE_CONFIG
// env variable, used only when options are nil.
func getConfigDefaults() (*Config, error) {
	fbc := &Config{}
	confFileName := os.Getenv(firebaseEnvName)
	if confFileName == "" {
		return fbc, nil
	}
	var dat []byte
	if confFileName[0] == byte('{') {
		dat = []byte(confFileName)
	} else {
		var err error
		if dat, err = ioutil.ReadFile(confFileName); err != nil {
			return nil, err
		}
	}
	if err := json.Unmarshal(dat, fbc); err != nil {
		return nil, err
	}

	// Some special handling necessary for db auth overrides
	var m map[string]interface{}
	if err := json.Unmarshal(dat, &m); err != nil {
		return nil, err
	}
	if ao, ok := m["databaseAuthVariableOverride"]; ok && ao == nil {
		// Auth overrides are explicitly set to null
		var nullMap map[string]interface{}
		fbc.AuthOverride = &nullMap
	}
	return fbc, nil
}

func getProjectID(ctx context.Context, config *Config, opts ...option.ClientOption) string {
	if config.ProjectID != "" {
		return config.ProjectID
	}

	creds, _ := transport.Creds(ctx, opts...)
	if creds != nil && creds.ProjectID != "" {
		return creds.ProjectID
	}

	if pid := os.Getenv("GOOGLE_CLOUD_PROJECT"); pid != "" {
		return pid
	}

	return os.Getenv("GCLOUD_PROJECT")
}
<|endoftext|>